Plotting Function Graphs Based on Formulas with Matplotlib and NumPy

Overview

This is a recipe for drawing mathematical function graphs in Python. Instead of entering point data manually, we use NumPy to generate an array for the domain (X-axis). By writing the formula directly into the code, you can efficiently draw smooth curves.

Specifications

• Input: Domain range (start, end, number of divisions) and the mathematical formula.
• Output: A visualized graph of the formula (displayed in a window or saved as an image).
• Requirements: matplotlib and numpy must be installed.

Basic Usage

import matplotlib.pyplot as plt
import numpy as np

# Create X-axis data from -10 to 10 divided into 100 points
x = np.linspace(-10, 10, 100)
# Calculate based on the formula (applied to all elements in the array)
y = x ** 2 + 5

fig, ax = plt.subplots()
ax.plot(x, y)
plt.show()

Full Code Example

The code below calculates a cubic function and a trigonometric function (sine wave) and plots them on the same graph for comparison.

import matplotlib.pyplot as plt
import numpy as np

def main():
    # 1. Generate the domain (X-axis)
    # Divide the range from -2π to 2π into 200 points
    x = np.linspace(-2 * np.pi, 2 * np.pi, 200)

    # 2. Calculate function values (Y-axis)
    # Operations between NumPy arrays are performed element-wise (broadcasting)
    
    # Function A: Cubic function (scaled by a coefficient)
    y1 = 0.1 * x ** 3 - x
    
    # Function B: Trigonometric function (sine wave)
    y2 = 3 * np.sin(x)

    # 3. Setup the graph
    fig, ax = plt.subplots(figsize=(8, 5))

    # 4. Plot and specify styles
    ax.plot(x, y1, label="Cubic Function", color="#1f77b4", linewidth=2)
    ax.plot(x, y2, label="Sine Wave", color="#ff7f0e", linestyle="--", linewidth=2)

    # 5. Decorate the graph
    ax.set_title("Function Comparison")
    ax.set_xlabel("x")
    ax.set_ylabel("f(x)")
    
    # Emphasize the axes (origin)
    ax.axhline(0, color='gray', linewidth=0.8)
    ax.axvline(0, color='gray', linewidth=0.8)
    
    ax.grid(True, linestyle=":", alpha=0.6)
    ax.legend() # Show the legend

    # Execute drawing
    plt.show()

if __name__ == "__main__":
    main()

Customization Points

• np.linspace(start, stop, num): This is the most important function for graph plotting. If the num (number of divisions) is too low, the line will look jagged. Use values like 100 or 200 to draw smooth curves.
• Axis Display (axhline, axvline): In mathematical graphs, it is often helpful to see the X and Y axes passing through the origin (0, 0). You can add them by specifying colors and thickness as shown in the code.
• Writing Formulas: NumPy includes major mathematical functions like np.sin(), np.cos(), np.exp() (exponential), and np.log() (logarithm). These functions accept arrays and return arrays.

Important Notes

• Confusion with the math library: The math.sin(x) function from the standard math library cannot handle lists or arrays directly and will cause an error. Always use NumPy functions like numpy.sin(x) when working with arrays for plotting.
• Values Outside the Domain: When plotting functions like np.log(x) or np.sqrt(x), if x contains negative values, it will produce nan (not a number) and those parts will not be drawn. You must adjust the linspace range to fit the function’s domain.
• Typos in Function Names: Be careful with function names in NumPy and Matplotlib as many are similar (e.g., use sin instead of siz, and linspace instead of linespace).

Variations (Optional)

Parametric Curves (Drawing a Circle)

By defining X and Y as independent functions, you can create graphs of parametric equations like circles or spirals.

import matplotlib.pyplot as plt
import numpy as np

def draw_circle():
    # Generate angle theta from 0 to 2π
    theta = np.linspace(0, 2 * np.pi, 100)
    
    # Equation for a circle: x = r * cos(θ), y = r * sin(θ)
    radius = 5
    x = radius * np.cos(theta)
    y = radius * np.sin(theta)

    fig, ax = plt.subplots()
    ax.plot(x, y)
    
    # Set aspect ratio to 1:1 so the circle is not distorted
    ax.set_aspect('equal')
    ax.grid(True)
    ax.set_title("Parametric Plot (Circle)")
    
    plt.show()

if __name__ == "__main__":
    draw_circle()

Summary

When creating graphs from formulas, using NumPy array operations allows you to write concise code without using loops. Once you understand the basic flow of setting the resolution with np.linspace and creating the drawing area with plt.subplots, you can visualize any mathematical function.