Scientific Computation MCP

from io import BytesIO from matplotlib.figure import Figure from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas from mcp.server.fastmcp import Image from sympy import symbols import sympy as sp from sympy.parsing.sympy_parser import standard_transformations, implicit_multiplication_application, parse_expr, \ convert_xor import numpy as np x, y, z = symbols("x y z") transforms = standard_transformations + (implicit_multiplication_application, convert_xor) local_ns = {"x": x, "y": y, "z": z, "sin": sp.sin, "cos": sp.cos} def register_tools(mcp): @mcp.tool() def plot_vector_field(f_str: str, bounds=(-1, 1, -1, 1, -1, 1), n: int = 10) -> Image: """ Plots a 3D vector field from a string "[u(x,y,z), v(x,y,z), w(x,y,z)]" Args: f_str: string representation of 3D field, e.g. "[z, -y, x]". bounds: (xmin, xmax, ymin, ymax, zmin, zmax) n: grid resolution per axis Returns: Displayed Matplotlib 3D quiver plot (no image return needed) """ # 1. Extract component strings raw = f_str.strip().lstrip("[").rstrip("]") u_s, v_s, w_s = [s.strip() for s in raw.split(",")] # 2. Parse each component with bare symbols u_expr = parse_expr(u_s, local_dict=local_ns, transformations=transforms) v_expr = parse_expr(v_s, local_dict=local_ns, transformations=transforms) w_expr = parse_expr(w_s, local_dict=local_ns, transformations=transforms) # 3. Convert to numpy functions u_fn = sp.lambdify((x, y, z), u_expr, "numpy") v_fn = sp.lambdify((x, y, z), v_expr, "numpy") w_fn = sp.lambdify((x, y, z), w_expr, "numpy") # 4. Prepare grid xmin, xmax, ymin, ymax, zmin, zmax = bounds X, Y, Z = np.meshgrid( np.linspace(xmin, xmax, n), np.linspace(ymin, ymax, n), np.linspace(zmin, zmax, n), indexing="ij" ) U = u_fn(X, Y, Z) V = v_fn(X, Y, Z) W = w_fn(X, Y, Z) try: fig = Figure(figsize=(8, 6)) canvas = FigureCanvas(fig) ax = fig.add_subplot(projection="3d") ax.quiver(X, Y, Z, U, V, W, length=0.1, normalize=True, color="blue") ax.set_xlabel("X") ax.set_ylabel("Y") ax.set_zlabel("Z") ax.set_title(f"3D Vector Field: {f_str}") # Save to buffer and return buf = BytesIO() canvas.print_png(buf) img_bytes = buf.getvalue() except ValueError as e: raise ValueError(f'Error plotting vector field: {e}') buf.close() return Image(data=img_bytes, format="png") @mcp.tool() def plot_function(expr_str: str, xlim: tuple[int, int] = (-5, 5), ylim: tuple[int, int] = (-5, 5), grid=200) \ -> Image: """ Plots a 2D or 3D mathematical function from a symbolic expression string. Args: expr_str: string representation of a function in x or x and y, e.g. "x**2" or "sin(sqrt(x**2 + y**2))" xlim: (xmin, xmax) range for x-axis ylim: (ymin, ymax) range for y-axis (used in 2D or 3D) grid: resolution of the plot grid Returns: A rendered Image of the function using Matplotlib. - 2D plot if the expression contains only x - 3D surface plot if the expression contains both x and y """ expr = parse_expr(expr_str, transformations=transforms, local_dict=local_ns) vars_used = expr.free_symbols # 2. Determine 2D or 3D if vars_used == {sp.symbols('x')} or vars_used == set(): # 2D: f(x) f_num = sp.lambdify(x, expr, modules=['numpy']) xs = np.linspace(xlim[0], xlim[1], grid) ys = f_num(xs) fig = Figure() ax = fig.add_subplot() ax.plot(xs, ys) ax.axhline(0, color='black', linewidth=0.8) ax.axvline(0, color='black', linewidth=0.8) ax.set_xlabel('x') ax.set_ylabel(expr_str) ax.set_title(f'f(x) = {expr_str}') elif vars_used >= {sp.symbols('x'), sp.symbols('y')}: # 3D: f(x, y) f_num = sp.lambdify((x, y), expr, modules=['numpy']) xs = np.linspace(xlim[0], xlim[1], grid) ys = np.linspace(ylim[0], ylim[1], grid) if ylim is not None else xs X, Y = np.meshgrid(xs, ys) Z = f_num(X, Y) fig = Figure() ax = fig.add_subplot(projection='3d') surf = ax.plot_surface(X, Y, Z, cmap='viridis') fig.colorbar(surf, ax=ax, shrink=0.6) ax.set_title(f'f(x, y) = {expr_str}') ax.set_xlabel('x') ax.set_ylabel('y') ax.set_zlabel(expr_str) else: raise ValueError("Plot only supports functions in x (2D) or x and y (3D).") # 3. Render as PNG buf = BytesIO() canvas = FigureCanvas(fig) canvas.print_png(buf) img_bytes = buf.getvalue() buf.close() return Image(data=img_bytes, format='png')