Graviton Python SDK

The Graviton SDK provides a seamless way to connect your local mathematical models (built with SymPy) to the Graviton Supercompute platform.

Installation

pip install graviton-sdk

Quick Start (Jupyter/Colab)

The SDK is designed to work beautifully in interactive environments.

import sympy as sp
import graviton as gv

# 1. Authenticate (Use the public evaluation key for prototyping)
gv.login('a-little-more-art-than-science')

# 2. Define a complex model
x, y = sp.symbols('x y')
f = sp.sin(x)**2 + sp.exp(-y) * sp.cos(x)

# 3. Display with the Graviton "Compute" button
gv.expr(f)

Programmatic Workflow

For advanced integrations, you can use the flattened API to evaluate and compute programmatically.

# 1. Evaluate to get JSON data
eval_data = gv.evaluate(f, options={"type": "JSON"})

# 2. Use flattened access for cost and ID
cost = eval_data['cost']
job_id = eval_data['id']

print(f"Estimated Cost: {cost:.2f} credits")

# 3. Compute using the ID directly
job = gv.compute(job_id, options={"auto_charge": True})
result = job.wait()
print(f"Result: {result}")

Astrophysics & N-Body Simulations (REBOUND & Mercury)

Graviton provides specialized Turing nodes for high-density N-body Hamiltonian simulations.

REBOUND Example

sim_payload = {
    "t": "rebound",
    "sim_time": 36500.0,
    "time_step": 0.1,
    "bodies": [
        {"name": "Sun", "mass": 1.0, "x": 0.0, "y": 0.0, "z": 0.0},
        {"name": "Earth", "mass": 3.0e-6, "a": 1.0, "e": 0.0167}
    ]
}

eval_data = gv.evaluate(sim_payload)
print(f"Estimated Cost: {eval_data.cost} credits")

job = gv.compute(eval_data.id)
results = job.wait()
print(results)

Mercury ETL Parser

If you have legacy Mercury parameter datasets (param.in, big.in, small.in), you can load them directly using gv.load_mercury:

# Automatically parses Mercury parameter files into a Graviton payload
mercury_payload = gv.load_mercury("./my_mercury_simulation/")

eval_data = gv.evaluate(mercury_payload)
job = gv.compute(eval_data.id)
results = job.wait()

Fluid Dynamics (WRF)

Graviton runs high-density WRF atmospheric modeling and weather simulations on Google Cloud Batch MPI clusters.

You can run WRF using three primary workflows:

1. Pre-Staged Datasets & Idealized Cases (Fastest, Zero-Upload)

Graviton provides pre-staged datasets (such as the standard NCAR CONUS 12km benchmark) and on-node Idealized Cases (generated via ideal.exe):

• Graviton-Hosted: "dataset": "conus_12km" (NCAR CONUS 12km HPC benchmark)
• Idealized Cases: "dataset": "ideal_b_wave" (3D Baroclinic Wave), "ideal_squall2d" (2D Squall Line), "ideal_quarter_ss" (3D Supercell), "ideal_tropical_cyclone" (3D Hurricane)

Because initial conditions are pre-staged in GCS or generated directly on the compute node, no file upload is required:

# Launch NCAR CONUS 12km benchmark directly
wrf_payload = {
    "t": "wrf",
    "dataset": "conus_12km",
    "inputs": {
        "sim_time": 3600  # Duration in seconds
    }
}
eval_data = gv.evaluate(wrf_payload)
print(f"Cost: {eval_data.cost:.2f} credits")

# Compute directly - no upload required!
job = gv.compute(eval_data.id)
results = job.wait()

2. Bring Your Own Storage (BYOS) & NOAA Buckets

Institutional researchers and climate labs can point Graviton directly to existing Google Cloud Storage buckets (or public NOAA datasets) using the storage (or s) parameter.

• NOAA Public Storage (Direct GCP Access):
  • NOAA HRRR 3km: "storage": "gs://high-resolution-rapid-refresh/" (or "dataset": "noaa_hrrr")
  • NOAA GFS 0.25°: "storage": "gs://global-forecast-system/" (or "dataset": "noaa_gfs")
• Institutional GCS Buckets:
  • Grant roles/storage.objectAdmin on your bucket to turing-machines@graviton-compute.iam.gserviceaccount.com.

wrf_byos_payload = {
    "t": "wrf",
    "storage": "gs://my-institution-bucket/hurricane-simulation/",
    "inputs": {
        "dx": 10000,
        "dy": 10000,
        "e_we": 150,
        "e_sn": 150,
        "sim_time": 86400
    }
}
eval_data = gv.evaluate(wrf_byos_payload)

# Cloud Batch streams wrfinput_d01 directly and writes wrfout files back to your bucket
job = gv.compute(eval_data.id)
results = job.wait()

For full details on dataset options, storage permissions, and cURL examples, visit the Graviton API Documentation.

3. Custom Upload Workflow (wrfinput_d01)

For custom modeling domains not hosted in cloud storage, evaluate to obtain a secure signed URL, upload your standard NetCDF wrfinput_d01 initial condition file, and trigger compute:

wrf_custom_payload = {
    "t": "wrf",
    "inputs": {
        "dx": 10000,
        "dy": 10000,
        "e_we": 100,
        "e_sn": 100,
        "sim_time": 43200
    }
}
eval_data = gv.evaluate(wrf_custom_payload)

# Upload the wrfinput_d01 NetCDF binary file
gv.load_WRF(eval_data.id, "path/to/wrfinput_d01")

# Launch compute
job = gv.compute(eval_data.id)
results = job.wait()

How it Works

When you call gv(formula), the SDK:

1. Serializes the SymPy object into a structured srepr format.
2. Generates a rich HTML representation with a LaTeX preview.
3. Adds a secure button that opens the Graviton Dashboard with your formula pre-loaded.

This allows you to leverage Graviton's massive cloud compute resources for complexity analysis and evaluation without ever leaving your research environment.

Security

Graviton follows a "Data, not Code" security model. The SDK only transmits the mathematical structure of your model, not arbitrary Python code. Your local environment remains isolated, and all authentication is handled securely via the Graviton web platform.