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

Learn how to calculate the trace (sum of diagonal elements) in NumPy using np.trace(), handle offset diagonals, and compute anti-diagonal sums for square and rectangular matrices.

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
])

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]
])

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

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]])

total = np.trace(matrix)

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]
])

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

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

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

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

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

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_diag = np.trace(np.fliplr(matrix))
print(f"Anti-diagonal sum: {anti_diag}")            # 15 (3 + 5 + 7)

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]])

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

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

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

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]
])

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

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

3D Arrays

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

import numpy as np

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

traces = np.trace(arr_3d, axis1=1, axis2=2)
print(f"Traces: {traces}")  # [5 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_via_trace = np.sqrt(np.trace(matrix.T @ matrix))
print(f"Via trace: {frobenius_via_trace:.4f}")

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.