from mpi4py import MPI
import numpy as np
import matplotlib.pyplot as plt

def create_jacobi_matrix(n, c1, c2, b):
    c = np.array([c1 if i % 2 == 0 else c2 for i in range(n)])
    b = np.full(n-1, b)
    return np.diag(c) + np.diag(b, k=1) + np.diag(b, k=-1)

def calculate_sturm_count(diagonal, off_diagonal, x):
    n = len(diagonal)
    p0 = 1.0
    p1 = diagonal[0] - x
    if np.isclose(p1, 0):
        p1 = 1e-10
    count = 0
    if p1 < 0:
        count += 1
    for i in range(1, n):
        p2 = (diagonal[i] - x) * p1 - off_diagonal[i-1]**2 * p0
        if np.isclose(p2, 0):
            p2 = 1e-10
        if p1 * p2 < 0:
            count += 1
        p0, p1 = p1, p2
    return count

# MPI initialization
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

# Parameters
n = 100
c1 = 12
c2 = 7
b = 15

# Create Jacobi matrix
jacobi_matrix = create_jacobi_matrix(n, c1, c2, b)
diagonal = np.diag(jacobi_matrix)
off_diagonal = np.diag(jacobi_matrix, k=1)

# Get eigenvalues
eigenvalues, _ = np.linalg.eig(jacobi_matrix)
interval_min, interval_max = eigenvalues.min(), eigenvalues.max()

# Create points for the range
num_points = 20  # Number of points
step = (interval_max - interval_min) / num_points
x_points = np.linspace(interval_min - step, interval_max + step, num_points + 2)

# Number of intervals
num_intervals = len(x_points) - 1

# Calculate intervals for each process
num_intervals_per_process = num_intervals // size
remainder = num_intervals % size

# Each process gets a chunk, with some getting one extra if there's a remainder
start_index = rank * num_intervals_per_process + min(rank, remainder)
end_index = start_index + num_intervals_per_process + (1 if rank < remainder else 0)

# Local computation for eigenvalue counts
local_eigenvalue_counts = np.array([calculate_sturm_count(diagonal, off_diagonal, x_points[i]) for i in range(start_index, end_index)], dtype=int)

# Gather results at the root process
counts = np.array([num_intervals_per_process + (1 if i < remainder else 0) for i in range(size)], dtype=int)
displacements = np.array([sum(counts[:i]) for i in range(size)], dtype=int)

all_eigenvalue_counts = None
if rank == 0:
    all_eigenvalue_counts = np.empty(num_intervals, dtype=int)

comm.Gatherv(sendbuf=local_eigenvalue_counts, recvbuf=(all_eigenvalue_counts, (counts, displacements)), root=0)

if rank == 0:
    # Calculate the differences correctly aligned
    count_differences_mpi = np.diff(all_eigenvalue_counts, prepend=0)

    # Sequential computation for comparison
    sequential_eigenvalue_counts = [calculate_sturm_count(diagonal, off_diagonal, x) for x in x_points]
    sequential_count_differences = np.diff(sequential_eigenvalue_counts, prepend=0)

    # Plot results
    plt.bar(x_points[:-1], sequential_count_differences, width=(x_points[1]-x_points[0]), align='edge', alpha=0.5, label='Sequential')
    plt.bar(x_points[:-1], -count_differences_mpi, width=(x_points[1]-x_points[0]), align='edge', alpha=0.5, label='MPI Computation')

    plt.xlabel('Intervals')
    plt.ylabel('Change in Eigenvalue Count')
    plt.title('Eigenvalue Counts per Interval')

    plt.xticks(ticks=x_points[:-1], labels=[f'{x:.2f}' for x in x_points[:-1]], rotation=45, ha="right")
    plt.legend()
    plt.tight_layout()
    plt.show()

    eigenvalues.sort()
    print("Eigenvalues:\n", eigenvalues)

Ошибка:

Traceback (most recent call last):
  File "C:\Users\home\Desktop\Simple\mpi2.py", line 86, in <module>
    plt.bar(x_points[:-1], sequential_count_differences, width=(x_points[1]-x_points[0]), align='edge', alpha=0.5, label='Sequential')
  File "C:\Users\home\AppData\Local\Programs\Python\Python312\Lib\site-packages\matplotlib\pyplot.py", line 2739, in bar
    return gca().bar(
           ^^^^^^^^^^
  File "C:\Users\home\AppData\Local\Programs\Python\Python312\Lib\site-packages\matplotlib\__init__.py", line 1478, in inner
    return func(ax, *map(sanitize_sequence, args), **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\home\AppData\Local\Programs\Python\Python312\Lib\site-packages\matplotlib\axes\_axes.py", line 2457, in bar
    x, height, width, y, linewidth, hatch = np.broadcast_arrays(
                                            ^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\home\AppData\Local\Programs\Python\Python312\Lib\site-packages\numpy\lib\stride_tricks.py", line 540, in broadcast_arrays
    shape = _broadcast_shape(*args)
            ^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\home\AppData\Local\Programs\Python\Python312\Lib\site-packages\numpy\lib\stride_tricks.py", line 422, in _broadcast_shape
    b = np.broadcast(*args[:32])
        ^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: shape mismatch: objects cannot be broadcast to a single shape.  Mismatch is between arg 0 with shape (21,) and arg 1 with shape (22,).