LLM Optimization Techniques for Scientific Computing

#engineering #oxlo #ai

Scientific computing scripts usually start as clean, naive implementations that fall over once the dataset grows. In this tutorial, I will build an LLM-powered optimization agent that ingests slow Python numerical code and returns vectorized, cache-friendly rewrites with complexity analysis. We will point it at Oxlo.ai so every optimization request costs one flat fee (see https://oxlo.ai/pricing), which makes the bill predictable even when we feed the model large simulation kernels.

What you'll need

Python 3.10 or newer
An Oxlo.ai API key from https://portal.oxlo.ai
The OpenAI SDK: pip install openai

Step 1: Set up the Oxlo.ai client and system prompt

I initialize the OpenAI-compatible client against Oxlo.ai and lock down a strict system prompt that forces the model to emit only a JSON optimization report. I use DeepSeek V3.2 here because it handles code restructuring well and sits on the free tier, so you can iterate without worrying about metered tokens.

from openai import OpenAI
import json

client = OpenAI(base_url="https://api.oxlo.ai/v1", api_key="YOUR_OXLO_API_KEY")

SYSTEM_PROMPT = """You are a scientific computing optimizer.
The user will provide Python code that performs numerical computation.
Analyze the code and produce an optimization report with exactly these fields:
- technique: the primary optimization applied (e.g., vectorization, cache blocking, algorithm substitution)
- complexity_before: Big-O notation of the original bottleneck
- complexity_after: Big-O notation of the optimized version
- optimized_code: a complete, runnable Python function that preserves the original API
- explanation: two sentences describing why the change helps

Return ONLY a JSON object matching this schema. Do not include markdown fences."""

Step 2: Build the optimization request handler

Next, I wrap the API call in a helper that sends the naive code and parses the JSON response. I keep the temperature low to reduce hallucinated imports.

def optimize_code(naive_code: str) -> dict:
    response = client.chat.completions.create(
        model="deepseek-v3.2",
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": naive_code},
        ],
        temperature=0.1,
        max_tokens=2048,
    )
    
    raw = response.choices[0].message.content.strip()
    if raw.startswith("

```json"):
        raw = raw[7:]
    if raw.startswith("```

"):
        raw = raw[3:]
    if raw.endswith("

```

"):
        raw = raw[:-3]
    return json.loads(raw.strip())

Step 3: Add a validation harness

Before trusting the rewrite, I verify that the optimized code is valid Python and that the original function name still exists. This catches syntax drift early.

import ast

def validate_optimized(original_name: str, optimized_code: str) -> bool:
    try:
        tree = ast.parse(optimized_code)
    except SyntaxError as e:
        print(f"Syntax error in optimized code: {e}")
        return False
    
    func_names = [
        node.name for node in ast.walk(tree)
        if isinstance(node, ast.FunctionDef)
    ]
    if original_name not in func_names:
        print(f"Warning: expected function '{original_name}' missing.")
        return False
    return True

Step 4: Prepare the naive workload

I define a triple-nested loop that computes pairwise Euclidean distances. This is the kind of bottleneck that appears in clustering and molecular dynamics pipelines, and it is an ideal candidate for vectorization.

NAIVE_CODE = '''
import numpy as np

def pairwise_distance(X):
    """X shape (n, d). Return (n, n) distance matrix."""
    n = X.shape[0]
    D = np.zeros((n, n))
    for i in range(n):
        for j in range(n):
            s = 0.0
            for k in range(X.shape[1]):
                diff = X[i, k] - X[j, k]
                s += diff * diff
            D[i, j] = np.sqrt(s)
    return D
'''

Run it

This script ties the pieces together. It submits the distance kernel to Oxlo.ai, prints the optimization report, validates the rewrite, then benchmarks both versions against random data to confirm the speedup is real.

import time
import numpy as np

# Generate report from Oxlo.ai
report = optimize_code(NAIVE_CODE)
print("Technique:", report.get("technique"))
print("Before:", report.get("complexity_before"))
print("After:", report.get("complexity_after"))
print("Explanation:", report.get("explanation"))
print("\n--- Optimized Code ---\n")
print(report["optimized_code"])

# Validate structure
assert validate_optimized("pairwise_distance", report["optimized_code"])

# Benchmark naive version
exec(NAIVE_CODE, globals())
X = np.random.rand(150, 40).astype(np.float64)
t0 = time.perf_counter()
out_naive = pairwise_distance(X)
t1 = time.perf_counter()
naive_sec = t1 - t0

# Benchmark optimized version
exec(report["optimized_code"], globals())
t0 = time.perf_counter()
out_opt = pairwise_distance(X)
t1 = time.perf_counter()
opt_sec = t1 - t0

# Verify numerical equivalence
np.testing.assert_allclose(out_naive, out_opt, rtol=1e-5)

print(f"\nNaive: {naive_sec:.4f}s")
print(f"Optimized: {opt_sec:.4f}s")
print(f"Speedup: {naive_sec / opt_sec:.1f}x")

Example output:

Technique: vectorization
Before: O(n^2 * d)
After: O(n^2 * d)
Explanation: Replaced Python loops with broadcasted NumPy operations to leverage SIMD and contiguous memory access. The asymptotic complexity is unchanged but the constant factor drops by orders of magnitude because the heavy work moves to C.

--- Optimized Code ---

import numpy as np

def pairwise_distance(X):
    """X shape (n, d). Return (n, n) distance matrix."""
    return np.sqrt(((X[:, None, :] - X[None, :, :]) ** 2).sum(axis=-1))

Naive: 1.2473s
Optimized: 0.0018s
Speedup: 692.9x

Wrap-up

You now have a working agent that rewrites slow numerical code using Oxlo.ai's flat per-request pricing. A concrete next step is to extend the prompt to request cache-blocking or Numba annotations when arrays exceed L3 cache. Another is to wire this into a pre-commit hook so every slow function committed to your repo gets an automatic optimization pass at a predictable cost.

DEV Community