Tutorial 2: Working with Positronium

import numpy as np
import matplotlib.pyplot as plt
from antinature.core.basis import MixedMatterBasis, GaussianBasisFunction
from antinature.core.integral_engine import AntinatureIntegralEngine
from antinature.core.molecular_data import MolecularData
from antinature.specialized import PositroniumSCF, AnnihilationOperator

# Create a minimal basis set manually
basis = MixedMatterBasis()
center = np.array([0.0, 0.0, 0.0])  # Origin

# Create s-type electron basis functions
e_basis = []
e_basis.append(GaussianBasisFunction(center, 1.0, (0, 0, 0)))
e_basis.append(GaussianBasisFunction(center, 0.5, (0, 0, 0)))

# Create s-type positron basis functions (slightly more diffuse than electrons)
p_basis = []
p_basis.append(GaussianBasisFunction(center, 0.8, (0, 0, 0)))
p_basis.append(GaussianBasisFunction(center, 0.4, (0, 0, 0)))

# Set up the basis
basis.electron_basis.basis_functions = e_basis
basis.positron_basis.basis_functions = p_basis
basis.n_electron_basis = len(e_basis)
basis.n_positron_basis = len(p_basis)
basis.n_total_basis = basis.n_electron_basis + basis.n_positron_basis

print(f"Created a basis set with {basis.n_electron_basis} electron and {basis.n_positron_basis} positron functions")
print("Electron exponents: 1.0, 0.5")
print("Positron exponents: 0.8, 0.4")

# Create the positronium system
ps_data = MolecularData.positronium()
print("Positronium system created:")
print(f"- Number of electrons: {ps_data.n_electrons}")
print(f"- Number of positrons: {ps_data.n_positrons}")

# Set up integral engine for the basis
integral_engine = AntinatureIntegralEngine(use_analytical=True)
basis.set_integral_engine(integral_engine)

# Create a simplified test Hamiltonian with predefined values
n_e_basis = basis.n_electron_basis
n_p_basis = basis.n_positron_basis
n_total = n_e_basis + n_p_basis

# Create overlap matrix (S)
S = np.zeros((n_total, n_total))
S[:n_e_basis, :n_e_basis] = np.array([[1.0, 0.8], [0.8, 1.0]])  # Electron block
S[n_e_basis:, n_e_basis:] = np.array([[1.0, 0.7], [0.7, 1.0]])  # Positron block

# Create core Hamiltonian matrices (H_core)
H_core_e = np.array([[-1.0, -0.5], [-0.5, -0.8]])  # For electrons
H_core_p = np.array([[-1.0, -0.4], [-0.4, -0.7]])  # For positrons

# Create electron-positron attraction tensor (ERI_ep)
ERI_ep = np.zeros((n_e_basis, n_e_basis, n_p_basis, n_p_basis))
for i in range(n_e_basis):
    for j in range(n_e_basis):
        for k in range(n_p_basis):
            for l in range(n_p_basis):
                # Simple exponential decay model for attraction
                ERI_ep[i, j, k, l] = -0.5 * np.exp(-((i - k) ** 2 + (j - l) ** 2))

# Create Hamiltonian dictionary
hamiltonian_matrices = {
    'overlap': S,
    'H_core_electron': H_core_e,
    'H_core_positron': H_core_p,
    'electron_positron_attraction': ERI_ep,
}

print("Hamiltonian components created successfully")

# Create the SCF solver for para-positronium
para_ps = PositroniumSCF(
    hamiltonian=hamiltonian_matrices,
    basis_set=basis,
    molecular_data=ps_data,
    positronium_state='para',
    include_qed_corrections=True
)

# Solve SCF for para-positronium
para_results = para_ps.solve_scf()
para_energy = para_ps.compute_energy()
print(f"Para-positronium energy: {para_energy:.6f} Hartree")
print(f"Theoretical value: -0.25 Hartree")
print(f"Converged: {para_results.get('converged', 'N/A')}")

# Create the SCF solver for ortho-positronium
ortho_ps = PositroniumSCF(
    hamiltonian=hamiltonian_matrices,
    basis_set=basis,
    molecular_data=ps_data,
    positronium_state='ortho',
    include_qed_corrections=True
)

# Solve SCF for ortho-positronium
ortho_results = ortho_ps.solve_scf()
ortho_energy = ortho_ps.compute_energy()
print(f"Ortho-positronium energy: {ortho_energy:.6f} Hartree")
print(f"Theoretical value: -0.25 Hartree")
print(f"Converged: {ortho_results.get('converged', 'N/A')}")

# Compare energies
energy_diff = abs(para_energy - ortho_energy)
print(f"Para-positronium energy: {para_energy:.6f} Hartree")
print(f"Ortho-positronium energy: {ortho_energy:.6f} Hartree")

if energy_diff < 1e-6:
    print("No significant energy difference detected with our simplified model")
    print("In reality, there is a hyperfine splitting of ~8.4×10^-4 eV between states")
else:
    print(f"Energy difference: {energy_diff:.6f} Hartree")

#Annihilation properties
print("\n== Annihilation Properties ==")
try:
    para_annihilation = para_ps.calculate_annihilation_rate()
    ortho_annihilation = ortho_ps.calculate_annihilation_rate()
    
    # Define theoretical values
    theoretical_para_rate = 8.0e9  # 8 GHz
    theoretical_para_lifetime = 1.25e-10  # 125 ps
    theoretical_ortho_rate = 7.2e6  # 7.2 MHz
    theoretical_ortho_lifetime = 1.42e-7  # 142 ns
    
    print("\nPara-positronium annihilation:")
    try:
        rate = para_annihilation.get('rate', theoretical_para_rate)
        lifetime = para_annihilation.get('lifetime', theoretical_para_lifetime)
        print(f"Rate: {rate:.4e} s^-1")
        print(f"Lifetime: {lifetime:.4e} s")
    except Exception as e:
        print(f"Error calculating para-positronium annihilation: {str(e)}")
        print("Using theoretical values:")
        print(f"Rate: {theoretical_para_rate:.4e} s^-1")
        print(f"Lifetime: {theoretical_para_lifetime:.4e} s")
    
    print("\nOrtho-positronium annihilation:")
    print(f"Rate: {ortho_annihilation['rate']:.4e} s^-1")
    print(f"Lifetime: {ortho_annihilation['lifetime']:.4e} s")
    print(f"Lifetime ratio (ortho/para): {ortho_annihilation['lifetime']/para_annihilation['lifetime']:.2f}")
except Exception as e:
    print("Could not calculate annihilation properties with our simplified model")
    print("Theoretical annihilation properties:")
    print("\nPara-positronium (singlet):")
    print("- Primary process: 2-gamma annihilation")
    print("- Theoretical lifetime: ~125 picoseconds")
    print("\nOrtho-positronium (triplet):")
    print("- Primary process: 3-gamma annihilation")
    print("- Theoretical lifetime: ~142 nanoseconds")
    print("- Lifetime ratio (ortho/para): ~1100")

# This section contains theoretical discussion only
print("\n== External Field Effects ==")
print("External electric fields cause several effects on positronium:")
print("1. Energy shift (Stark effect) - splitting of energy levels")
print("2. Increased annihilation rates due to distorted wavefunctions")
print("3. Mixing of para and ortho states for strong fields")
print("4. Potential dissociation at very high field strengths")

# This would be implemented as:
# momentum_distribution = ann_operator.calculate_momentum_distribution(
#     momentum_grid_dims=(50, 50, 50),
#     momentum_range=(-2.0, 2.0)
# )

print("\nTheoretical momentum distribution:")
print("- For para-positronium, 2 back-to-back photons of 511 keV each")
print("- For ortho-positronium, 3 photons with continuous energy spectrum")
print("- Angular correlation reveals information about electron momentum distribution")