Python NumPy: How to Sum Diagonal Elements in NumPy in Python

The sum of diagonal elements (called the trace) is a fundamental matrix property used in linear algebra, eigenvalue calculations, and matrix analysis. NumPy provides dedicated functions for this operation.

Using np.trace() - Recommended

The trace() function directly returns the sum of diagonal elements:

import numpy as np

matrix = np.array([
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9]
])

trace = np.trace(matrix)
print(trace)  # 15 (1 + 5 + 9)

Visual Representation

import numpy as np

matrix = np.array([
    [1, 2, 3],    # ↙ 1 is on main diagonal
    [4, 5, 6],    #   ↙ 5 is on main diagonal
    [7, 8, 9]     #     ↙ 9 is on main diagonal
])

# Main diagonal: positions (0,0), (1,1), (2,2)
print(np.trace(matrix))  # 1 + 5 + 9 = 15

Using np.diagonal()

When you need the diagonal elements themselves (not just the sum):

import numpy as np

matrix = np.array([
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9]
])

# Get diagonal elements as array
diag_elements = np.diagonal(matrix)
print(f"Diagonal: {diag_elements}")     # [1 5 9]

# Then sum if needed
print(f"Sum: {np.sum(diag_elements)}")  # 15

Output:

Diagonal: [1 5 9]
Sum: 15

When to Use Each

import numpy as np

matrix = np.array([[1, 2], [3, 4]])

# Just need the sum? Use trace()
total = np.trace(matrix)

# Need to process elements? Use diagonal()
diag = np.diagonal(matrix)
squared = diag ** 2
mean = diag.mean()

print(total)    # 5
print(diag)     # [1 4]
print(squared)  # [ 1 16]
print(mean)     # 2.5

Offset Diagonals

Access diagonals above or below the main diagonal using the offset parameter:

import numpy as np

matrix = np.array([
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9]
])

# Main diagonal (offset=0, default)
print(f"Main (k=0): {np.trace(matrix, offset=0)}")     # 15

# One above main diagonal (offset=1)
print(f"Upper (k=1): {np.trace(matrix, offset=1)}")    # 8 (2 + 6)

# One below main diagonal (offset=-1)
print(f"Lower (k=-1): {np.trace(matrix, offset=-1)}")  # 12 (4 + 8)

# Two above main diagonal (offset=2)
print(f"Upper (k=2): {np.trace(matrix, offset=2)}")    # 3

Output:

Main (k=0): 15
Upper (k=1): 8
Lower (k=-1): 12
Upper (k=2): 3

Visualizing Offset Diagonals

Matrix:
    [1, 2, 3]
    [4, 5, 6]
    [7, 8, 9]

offset=2:  [3]                    → sum = 3
offset=1:  [2, 6]                 → sum = 8
offset=0:  [1, 5, 9]  (main)      → sum = 15
offset=-1: [4, 8]                 → sum = 12
offset=-2: [7]                    → sum = 7

Non-Square Matrices

Trace works on rectangular matrices too:

import numpy as np

# More columns than rows
wide = np.array([
    [1, 2, 3, 4],
    [5, 6, 7, 8]
])
print(f"Wide matrix trace: {np.trace(wide)}")  # 7 (1 + 6)

# More rows than columns
tall = np.array([
    [1, 2],
    [3, 4],
    [5, 6]
])
print(f"Tall matrix trace: {np.trace(tall)}")  # 5 (1 + 4)

Output:

Wide matrix trace: 7
Tall matrix trace: 5

Anti-Diagonal Sum

The anti-diagonal runs from top-right to bottom-left:

import numpy as np

matrix = np.array([
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9]
])

# Anti-diagonal: flip horizontally, then get trace
anti_diag = np.trace(np.fliplr(matrix))
print(f"Anti-diagonal sum: {anti_diag}")            # 15 (3 + 5 + 7)

# Or get elements using diagonal on flipped matrix
anti_elements = np.diagonal(np.fliplr(matrix))
print(f"Anti-diagonal elements: {anti_elements}")   # [3 5 7]

Output:

Anti-diagonal sum: 15
Anti-diagonal elements: [3 5 7]

Trace Properties

The trace has useful mathematical properties:

import numpy as np

A = np.array([[1, 2], [3, 4]])
B = np.array([[5, 6], [7, 8]])

# Property 1: tr(A + B) = tr(A) + tr(B)
print(f"tr(A+B) = {np.trace(A + B)}")                  # 18
print(f"tr(A) + tr(B) = {np.trace(A) + np.trace(B)}")  # 18

# Property 2: tr(cA) = c * tr(A)
c = 3
print(f"tr(3A) = {np.trace(c * A)}")                    # 15
print(f"3 * tr(A) = {c * np.trace(A)}")                 # 15

# Property 3: tr(A^T) = tr(A)
print(f"tr(A) = {np.trace(A)}")                         # 5
print(f"tr(A.T) = {np.trace(A.T)}")                     # 5

# Property 4: tr(AB) = tr(BA)
print(f"tr(AB) = {np.trace(A @ B)}")                    # 69
print(f"tr(BA) = {np.trace(B @ A)}")                    # 69

Output:

tr(A+B) = 18
tr(A) + tr(B) = 18
tr(3A) = 15
3 * tr(A) = 15
tr(A) = 5
tr(A.T) = 5
tr(AB) = 69
tr(BA) = 69

Sum of Eigenvalues

The trace equals the sum of eigenvalues:

import numpy as np

matrix = np.array([
    [4, 2],
    [1, 3]
])

# Calculate trace
trace = np.trace(matrix)
print(f"Trace: {trace}")                                # 7

# Calculate eigenvalues
eigenvalues = np.linalg.eigvals(matrix)
print(f"Eigenvalues: {eigenvalues}")                    # [5. 2.]
print(f"Sum of eigenvalues: {np.sum(eigenvalues)}")     # 7.0

Output:

Trace: 7
Eigenvalues: [5. 2.]
Sum of eigenvalues: 7.0

note

The trace equals the sum of eigenvalues, and the determinant equals their product. These relationships are useful for checking calculations.

3D Arrays

For 3D arrays, trace operates on the last two dimensions by default:

import numpy as np

# Stack of 2x2 matrices
arr_3d = np.array([
    [[1, 2], [3, 4]],
    [[5, 6], [7, 8]]
])

# Trace of each 2D matrix
traces = np.trace(arr_3d, axis1=1, axis2=2)
print(f"Traces: {traces}")  # [5 13]
# First matrix: 1+4=5, Second matrix: 5+8=13

Output:

Traces: [ 5 13]

Practical Example: Matrix Norm

The Frobenius norm can be computed using trace:

import numpy as np

matrix = np.array([
    [1, 2],
    [3, 4]
])

# Frobenius norm = sqrt(trace(A^T @ A))
frobenius_via_trace = np.sqrt(np.trace(matrix.T @ matrix))
print(f"Via trace: {frobenius_via_trace:.4f}")

# Compare with np.linalg.norm
frobenius_direct = np.linalg.norm(matrix, 'fro')
print(f"Direct: {frobenius_direct:.4f}")

Output:

Via trace: 5.4772
Direct: 5.4772

Summary

Method	Returns	Use Case
np.trace(A)	Scalar (sum)	Direct diagonal sum
np.diagonal(A)	Array (elements)	Need individual values
np.trace(A, offset=k)	Scalar	Sum of offset diagonal
np.trace(np.fliplr(A))	Scalar	Anti-diagonal sum

Use np.trace() when you only need the sum of diagonal elements-it's more efficient than extracting the diagonal first. Use np.diagonal() when you need to process the individual diagonal values.