Here is my discrete ordinates transport code that I generated for fun, problem is that the output for a flash of light isn't correct. Are you able to tell me what is wrong with the code?     import numpy as np
import matplotlib.pyplot as plt

# Constants and Parameters
# ------------------------
Nx, Ny = 37, 37  # Mesh size
G = 1             # Single energy group for simplicity
Nm = 4            # Number of discrete ordinates (e.g., S2 quadrature)
Nm = 8            # Number of discrete ordinates (e.g., S4 quadrature)
Lx, Ly = 10.0, 10.0  # Mesh dimensions in cm
Sigma_t = 0.5     # Total cross-section (cm^-1)
Sigma_s = 0.3     # Scattering cross-section (cm^-1)
Sigma_a = Sigma_t - Sigma_s  # Absorption cross-section (cm^-1)
Q = 1.0           # External Source (cm^-3 s^-1)
Emission_Fraction = 0.1  # Fraction of internal energy emitted per iteration
BCs = 'vacuum'     # Boundary condition type

# Discrete Ordinates Directions (Simplified S2 Quadrature)
mu = np.array([0.57735, 0.57735, -0.57735, -0.57735])
eta = np.array([0.57735, -0.57735, 0.57735, -0.57735])
w = np.array([0.34, 0.34, 0.34, 0.551])
# # Discrete Ordinates Directions and Weights for S4 Quadrature
mu = np.array([
    0.2958759478974849,  0.2958759478974849, -0.2958759478974849, -0.2958759478974849,
    0.9081970464923464, -0.9081970464923464,  0.2958759478974849, -0.2958759478974849
])
eta = np.array([
    0.9081970464923464, -0.9081970464923464,  0.2958759478974849, -0.2958759478974849,
    0.2958759478974849,  0.2958759478974849, -0.2958759478974849, -0.2958759478974849
])
w = np.array([
    0.585786438095848,  0.585786438095848,  0.585786438095848,  0.585786438095848,
    0.121518570687221,  0.121518570687221,  0.121518570687221,  0.121518570687221
])

# Initialize Photon Fluxes and Cell Internal Energies
psi = np.zeros((G, Nm, Ny, Nx))
internal_energy = np.zeros((Ny, Nx))

# Set Initial Source
psi[:, :, Ny//2, Nx//2] = Q / Sigma_t

# Main Simulation Loop
for _ in range(30):  
    psi_new = np.copy(psi)
    internal_energy_new = np.copy(internal_energy)
    
    # Sweep over all directions and cells
    for g in range(G):
        for m in range(Nm):
            for i in range(Nx):
                for j in range(Ny):
                    # Calculate the flux for this direction and cell
                    psi_out = psi[g, m, j, i]
                    
                    # Energy Deposition
                    energy_deposited = Sigma_a * psi_out
                    internal_energy_new[j, i] += energy_deposited
                    
                    # Update flux for the next iteration based on the direction
                    if mu[m] > 0 and i < Nx - 1:
                        psi_new[g, m, j, i+1] += psi_out*abs(mu[m])*w[m]
                    elif mu[m] < 0 and i > 0:
                        psi_new[g, m, j, i-1] += psi_out*abs(mu[m])*w[m]
                    
                    if eta[m] > 0 and j < Ny - 1:
                        psi_new[g, m, j+1, i] += psi_out*abs(eta[m])*w[m]
                    elif eta[m] < 0 and j > 0:
                        psi_new[g, m, j-1, i] += psi_out*abs(eta[m])*w[m]
    
    # Photon Emission from Internal Energy
    emission = Emission_Fraction * internal_energy
    internal_energy_new -= emission  # Update internal energy after emission

    for g in range(G):
        for j in range(Ny):
            for i in range(Nx):
                psi_new[g, :, j, i] += emission[j, i] / Nm  # Isotropic emission distributed across all directions

    
    # Apply Boundary Conditions
    if BCs == 'vacuum':
        psi_new[:, :, 0, :] = 0  
        psi_new[:, :, -1, :] = 0  
        psi_new[:, :, :, 0] = 0  
        psi_new[:, :, :, -1] = 0  
    
    psi, internal_energy = psi_new, internal_energy_new

# Plot the Final Scalar Flux and Internal Energy
scalar_flux = np.sum(psi, axis=(0, 1))
plt.figure(figsize=(12, 5))

plt.subplot(1, 2, 1)
plt.imshow(scalar_flux, cmap='viridis', origin='lower')
plt.title("Final Scalar Flux")
plt.colorbar(label='Flux (arb. units)')

plt.subplot(1, 2, 2)
plt.imshow(internal_energy, cmap='plasma', origin='lower')
plt.title("Final Internal Energy Distribution")
plt.colorbar(label='Energy (arb. units)')

plt.tight_layout()
plt.show()