Enabling Your GPU for a Solver in MusicBox

This tutorial will show you to use 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.9
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 largerly the same, but with an additional is_cuda_available() function to verify that the GPU is running properly:

from acom_music_box import MusicBox, Conditions
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 outside the main code 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 reactions that the species undergo in the
    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
    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.initial_conditions = Conditions(
        temperature=298.15, # Units: Kelvin (K)
        pressure=101325.0, # Units: Pascals (Pa)
        species_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.add_evolving_condition(
        100.0, # Units: Seconds (s)
        Conditions(
            temperature=75.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