Transforming spin-orbit coupling from the total angular momentum base to m1 m2

cg.py

import sympy as sp
import sympy.physics.quantum as spq
from sympy.matrices.dense import matrix_multiply_elementwise as mme

import numpy as np

import itertools
from fractions import Fraction
from argparse import ArgumentParser

# # # # # # # # # # # # # # # # #  # #
# Angular Momentum Physics Functions #
# # # # # # # # # # # # # # # # #  # #


def generate_M1M2_base(J1, J2):
    """
    Generate the angular momemtum space in the |m1 m2> base.
    
    The base is generated in an order such that the transformation matrix
    between |m1 m2> and |jtot m> is of block diagonal form.
    """
    Mmax = abs(J1 + J2)

    M1_values = [J1 - i for i in range(2 * J1 + 1)]
    M2_values = [J2 - i for i in range(2 * J2 + 1)]
    M_values = [Mmax - i for i in range(2 * Mmax + 1)]

    M1M2_combos = list(itertools.product(M1_values, M2_values))

    M1M2_base = [
        [
            spq.Ket(*valid_combo)
            for valid_combo in filter(lambda x: (x[0] + x[1]) == M, M1M2_combos)
        ]
        for M in M_values
    ]

    M1M2_base = list(itertools.chain.from_iterable(M1M2_base))

    return M1M2_base


def generate_JM_base(J1, J2):
    """
    Generate the angular momemtum space in the |jtot m> base.
    
    The base is generated in an order such that the transformation matrix
    between |m1 m2> and |jtot m> is of block diagonal form.
    """
    Mmax = abs(J1 + J2)
    Jmax = abs(J1 + J2)
    Jmin = abs(J1 - J2)

    M_values = [Mmax - i for i in range(2 * Mmax + 1)]
    J_values = [Jmax - i for i in range(Jmax - Jmin + 1)]

    JM_base = [spq.Ket(*(J, M)) for (M, J) in itertools.product(M_values, J_values) if abs(M) <= J]

    return JM_base


def transformation_matrix(jm_base, m1m2_base):
    """
    Create transformation matrix between |jtot m> and |m1 m2>.
    """

    # J1 and J2 can be deduced from the max value that m1 and m2 take.
    J1 = max([base_ket.label[0] for base_ket in m1m2_base])
    J2 = max([base_ket.label[1] for base_ket in m1m2_base])

    # Special CG for our case of interest with J1 and J2 fixed.
    def cg_(m1m2_, jm_):
        coeff = spq.cg.CG(
            J1, m1m2_.label[0], J2, m1m2_.label[1], jm_.label[0], jm_.label[1]
        ).doit()
        return coeff

    U = sp.Matrix(
        [
            [cg_(m1m2_base[ii], jm_base[jj]) for jj in range(len(jm_base))]
            for ii in range(len(m1m2_base))
        ]
    )

    return spq.Dagger(U)


def LS_jtot_base(jm_base):
    """
    LS = J**2 - L**2 - S**2
    Diagonal in the total angular momentum base. Note that this eq
    doesn't depend on the exchange of L and S.
    """
    Jmax = max([base_ket.label[0] for base_ket in jm_base])
    Jmin = min([base_ket.label[0] for base_ket in jm_base])

    # The order of J1/J2 may be wrong do to an ambiguity with an
    # abs value but we don't care about it here.
    J1 = (Jmax + Jmin) / 2
    J2 = (Jmax - Jmin) / 2

    diag_elements_LS_ = lambda j: (j * (j + 1) - J1 * (J1 + 1) - J2 * (J2 + 1))
    diagonal_LS = [
        diag_elements_LS_(base_vector.label[0]) for base_vector in jm_base
    ]
    LS = sp.diag(*diagonal_LS) / 2

    return LS


def change_basis(transformation_matrix, operator):
    return spq.Dagger(transformation_matrix) * operator * transformation_matrix


# Extra misc tools #
def parse_ratio_to_symbol(string_ratio):
    return sp.S(Fraction(string_ratio))


def insert_in_template(template, replacements):
    for replacement in replacements:
        template = template.replace(
            replacement, sp.latex(replacements[replacement], mode="equation")
        )
    return template


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("j1", type=parse_ratio_to_symbol)
    parser.add_argument("j2", type=parse_ratio_to_symbol)
    parser.add_argument("--no-terminal", action="store_true")
    parser.add_argument("--latex", action="store_true")
    args = parser.parse_args()

    # Check inputs
    J1 = abs(args.j1)
    J2 = abs(args.j2)

    if not ((J1 * 2).is_integer and (J2 * 2).is_integer):
        raise ValueError("Spins have to be multiples of 1/2!")

    # Do some angular momentum physics
    jm_base = generate_JM_base(J1, J2)
    m1m2_base = generate_M1M2_base(J1, J2)

    V = transformation_matrix(jm_base, m1m2_base)
    LS = LS_jtot_base(jm_base)
    LS_m1m2_basis = change_basis(V, LS)

    # Print results
    replacements = {
        "@LS_JM": LS,
        "@LS_M1M2": LS_m1m2_basis,
        "@V": mme(np.sign(V), mme(V, V)),
        "@SPACE_M1M2": sp.Matrix(m1m2_base),
        "@SPACE_JM": sp.Matrix(jm_base),
    }

    if not args.no_terminal:
        sp.init_printing()
        print("\nTransformation matrix from |jtot m> to |m1 m2>")
        print("This is not the actual matrix; elements are represented as in the")
        print("CG tables. Do element-wise sqrt but keep the sign for the matrix.")
        sp.pprint(replacements["@V"])
        print("\nLS in the |jtot m> basis")
        sp.pprint(replacements["@LS_JM"])
        print("\nLS in the |m1 m2> basis")
        sp.pprint(replacements["@LS_M1M2"])
        print("\nBasis |jtot m>")
        sp.pprint(replacements["@SPACE_M1M2"])
        print("\nBasis |m1 m2>")
        sp.pprint(replacements["@SPACE_JM"])

    if args.latex:
        print("Generating LaTex!")
        template_filename = "./latex_template.tex"
        out_filename = "outfile.tex"

        with open(template_filename, "r") as template:
            data = "".join(template.readlines())

        with open(out_filename, "w") as outfile:
            outfile.write(insert_in_template(data, replacements))

Usage:
python3 j1 j2 --no-terminal --latex

Example 1. No latex output and terminal output:
python3 3/2 1/2

Example 2. LaTex output and terminal output:
python3 3/2 1 --latex

Example 3. LaTex output and NO terminal output:
python3 3/2 1/2 --latex --no-terminal

Example 4. No latex and No terminal output (Why would you want this?)
python3 3/2 1/3 --no-terminal # This will raise an exception due to the j2=1/3 spin.

Notes:

• The --no-terminal flag is optional. It disables the output to standard output.
• The --latex flag is optional. For it to work you will also need a latex template which can be found on this other gist. A output.tex file with the output will be generated which you can then compile with pdflatex output.tex. Pathnames are hard-coded so the template file will have to be in the same directory as this script and named latex_template.tex.
• j1 and j2 are the spin of the subspace written as ratios (e.g. 3/2).
• If you try to give a value for j1 or j2 that is not a multiple of 1/2 a ValueError exception will be raised.
• Negative values for j1 and j2 will be made positive automatically via abs.