plot.py

import matplotlib.pyplot as plt
import numpy as np
import streamlit as st
import sympy as sp

import model


def base_plot_setup():
    fig, ax = plt.subplots(figsize=(10, 5))
    ax.axhline(0, color='black', lw=1.2)
    ax.axvline(0, color='black', lw=1.2)
    ax.grid(True, linestyle='--', alpha=0.6)
    return fig, ax


def finalize_plot(fig, ax):
    ax.legend(loc="upper right", bbox_to_anchor=(1, 1), frameon=True, shadow=True)
    st.pyplot(fig)


def plot_single(f_num, root, a, b, func_label):
    fig, ax = base_plot_setup()
    x_range = np.linspace(a - 2, b + 2, 500)
    y_range = f_num(x_range)

    ax.plot(x_range, y_range, label=f"f(x) = {func_label}", color='#1f77b4', lw=2)

    if root is not None:
        ax.scatter([root], [f_num(root)], color='red', zorder=5, label=f"Корень")
        ax.annotate(f"{root:.3f}", (root, 0), textcoords="offset points", xytext=(0, 10), ha='center')

    ax.set_xlim(a - 1, b + 1)
    finalize_plot(fig, ax)


def plot_system(f_list, solution, point):
    fig, ax = base_plot_setup()
    offset = 3

    nodes = [point[0], point[1]]
    limit = max(map(abs, nodes)) + offset

    x = np.linspace(-limit, limit, 100)
    y = np.linspace(-limit, limit, 100)
    X, Y = np.meshgrid(x, y)

    for i, f in enumerate(f_list):
        Z = f(X, Y)
        cs = ax.contour(X, Y, Z, levels=[0], colors=['#1f77b4', '#ff7f0e'][i])
        cs.set_label(f"Уравнение {i + 1}: {f}")

    ax.scatter([point[0]], [point[1]], color='gray', label='Начало', zorder=5)
    if solution is not None:
        ax.scatter([solution[0]], [solution[1]], color='red', label='Решение', zorder=6)

    ax.set_xlabel("x")
    ax.set_ylabel("y")
    finalize_plot(fig, ax)