Enabling Your GPU for a Solver in MusicBox

This tutorial will show you how to utilize a GPU for your MusicBox work.

However, this tutorial will not cover how to efficiently use a GPU through parallelization; it will simply introduce getting a GPU set up to run your code.

Note: This tutorial requires you to have a Linux GPU-ready environment handy, such as a supercomputing node; it will fail otherwise.

1. Creating a GPU Virtual Environment

Running code on a GPU requires a different install protocol when setting up a virtual environment.

To do so, run these commands in your terminal:

conda create --name music_box_gpu python=3.13
conda activate music_box_gpu
pip install --upgrade setuptools pip wheel
pip install nvidia-pyindex
pip install acom_music_box
pip install musica[gpu]
conda install ipykernel scikit-learn seaborn scipy dask

2. Importing MusicBox

Importing MusicBox is largely the same, but with an additional is_cuda_available() function to verify that the GPU is running properly:

from acom_music_box import MusicBox
import musica.mechanism_configuration as mc
import matplotlib.pyplot as plt
from musica.cuda import is_cuda_available
from musica import SolverType

As with creating the music_box environment in the Basic Workflow Tutorial, this cell may be slow to run the first time.

3. Running a Basic Solver on GPU

This code is a copy of the Basic Workflow Tutorial, but with an if statement added to verify that it is running on a GPU.

If you are seeing “Error: No GPU Available” being printed, that means a GPU was not detected; verify that your environment has a GPU.

if is_cuda_available():
    # Create each of the species that will be simulated
    X = mc.Species(name="X")
    Y = mc.Species(name="Y")
    Z = mc.Species(name="Z")
    species = {"X": X, "Y": Y, "Z": Z}
    gas = mc.Phase(name="gas", species=list(species.values()))

    # Create the Arrhenius reactions
    arr1 = mc.Arrhenius(name="X->Y", A=4.0e-3, C=50, reactants=[species["X"]], products=[species["Y"]], gas_phase=gas)
    arr2 = mc.Arrhenius(name="Y->Z", A=4.0e-3, C=50, reactants=[species["Y"]], products=[species["Z"]], gas_phase=gas)
    rxns = {"X->Y": arr1, "Y->Z": arr2}

    # Create the mechanism that is defined by the species, phases, and reactions
    mechanism = mc.Mechanism(name="tutorial_mechanism", species=list(species.values()), phases=[gas], reactions=list(rxns.values()))

    # Create the box model that contains the mechanism with GPU solver
    box_model = MusicBox()
    box_model.load_mechanism(mechanism, solver_type=SolverType.cuda_rosenbrock)

    # Set the conditions of the box model at time = 0 s
    box_model.set_condition(
        time=0,
        temperature=298.15,  # Units: Kelvin (K)
        pressure=101325.0,   # Units: Pascals (Pa)
        concentrations={     # Units: mol/m^3
            "X": 3.75,
            "Y": 5.0,
            "Z": 2.5,
        }
    )

    # Set the box model conditions at the defined time
    box_model.set_condition(
        time=100.0,         # Units: Seconds (s)
        temperature=275.0,  # Units: Kelvin (K)
        pressure=100100.0   # Units: Pascals (Pa)
    )

    # Set the additional configuration options for the box model
    box_model.box_model_options.simulation_length = 200  # Units: Seconds (s)
    box_model.box_model_options.chem_step_time = 1       # Units: Seconds (s)
    box_model.box_model_options.output_step_time = 20    # Units: Seconds (s)

    df = box_model.solve()
    display(df)
    df.plot(x='time.s', y=['CONC.X.mol m-3', 'CONC.Y.mol m-3', 'CONC.Z.mol m-3'],
            title='Concentration over time',
            ylabel='Concentration (mol m-3)',
            xlabel='Time (s)')
    plt.show()
else:
    print("Error: No GPU Available")

Error: No GPU Available