Profiling calculation¶

It can be interesting to know where a calculation is spending its time. To this end, rascaline includes self-profiling code that can record and display which part of the calculation takes time, and which function called long-running functions. All the example should output something similar to the table below.

╔════╦══════════════════════════════╦════════════╦═══════════╦══════════╦═════════════╗
║ id ║ span name                    ║ call count ║ called by ║ total    ║ mean        ║
╠════╬══════════════════════════════╬════════════╬═══════════╬══════════╬═════════════╣
║  2 ║ Full calculation             ║          1 ║         — ║ 660.58ms ║    660.58ms ║
╠════╬══════════════════════════════╬════════════╬═══════════╬══════════╬═════════════╣
║  3 ║ SoapPowerSpectrum::compute   ║          1 ║         2 ║ 584.02ms ║    584.02ms ║
╠════╬══════════════════════════════╬════════════╬═══════════╬══════════╬═════════════╣
║  1 ║ Calculator::prepare          ║          2 ║      3, 2 ║ 148.15ms ║     74.08ms ║
╠════╬══════════════════════════════╬════════════╬═══════════╬══════════╬═════════════╣
║  0 ║ NeighborsList                ║         20 ║         1 ║  20.82ms ║      1.04ms ║
╠════╬══════════════════════════════╬════════════╬═══════════╬══════════╬═════════════╣
║  5 ║ SphericalExpansion::compute  ║          1 ║         3 ║ 196.38ms ║    196.38ms ║
╠════╬══════════════════════════════╬════════════╬═══════════╬══════════╬═════════════╣
║  4 ║ GtoRadialIntegral::compute   ║      74448 ║         5 ║ 117.04ms ║      1.57µs ║
╠════╬══════════════════════════════╬════════════╬═══════════╬══════════╬═════════════╣
║  6 ║ SphericalHarmonics::compute  ║      74448 ║         5 ║   9.95ms ║ 133.00ns ⚠️ ║
╚════╩══════════════════════════════╩════════════╩═══════════╩══════════╩═════════════╝

In this table, the first columns assign a unique numeric identifier to each section of the code. The second one displays the name of the section. Then come the number of time this section of the code have been executed, which other function/section called the current one, and finally the total and mean time spent in this function.

The ⚠️ symbol is added when the mean cost of the function is close to the profiling overhead (30 to 80ns per function call), and thus the measurement might not be very reliable.

Some of the most important sections are:

Calculator::prepare: building the list of samples/properties that will be in the descriptor
XXX::compute: building blocks for the overall calculation
NeighborsList: construction of the list of neighbors

You can obtain a dataset for profiling from our website.

import chemfiles

import rascaline
from rascaline import SoapPowerSpectrum


def compute_soap(path):
    """Compute SOAP power spectrum.

    This is the same code as the 'compute-soap' example
    """
    with chemfiles.Trajectory(path) as trajectory:
        frames = [f for f in trajectory]

    HYPER_PARAMETERS = {
        "cutoff": 5.0,
        "max_radial": 6,
        "max_angular": 4,
        "atomic_gaussian_width": 0.3,
        "center_atom_weight": 1.0,
        "radial_basis": {
            "Gto": {},
        },
        "cutoff_function": {
            "ShiftedCosine": {"width": 0.5},
        },
    }

    calculator = SoapPowerSpectrum(**HYPER_PARAMETERS)
    descriptor = calculator.compute(frames, gradients=["positions"])
    descriptor = descriptor.keys_to_samples("center_type")
    descriptor = descriptor.keys_to_properties(["neighbor_1_type", "neighbor_2_type"])

    return descriptor

Run the calculation with profiling enabled.

with rascaline.Profiler() as profiler:
    descriptor = compute_soap("dataset.xyz")

Display the recorded profiling data as table.

print(profiler.as_short_table())

╔════╦══════════════════════════════════════════════╦════════════╦═══════════╦══════════╦══════════════╗
║ id ║ span name                                    ║ call count ║ called by ║ total    ║ mean         ║
╠════╬══════════════════════════════════════════════╬════════════╬═══════════╬══════════╬══════════════╣
║  4 ║ SoapPowerSpectrum::compute                   ║          1 ║         — ║ 420.09ms ║     420.09ms ║
╠════╬══════════════════════════════════════════════╬════════════╬═══════════╬══════════╬══════════════╣
║  5 ║ SphericalExpansion::compute                  ║          1 ║         4 ║ 197.61ms ║     197.61ms ║
╠════╬══════════════════════════════════════════════╬════════════╬═══════════╬══════════╬══════════════╣
║  1 ║ SoapRadialIntegralSpline::with_accuracy      ║          6 ║         5 ║  53.03ms ║       8.84ms ║
╠════╬══════════════════════════════════════════════╬════════════╬═══════════╬══════════╬══════════════╣
║  0 ║ GtoRadialIntegral::compute                   ║       3846 ║         1 ║  50.16ms ║      13.04µs ║
╠════╬══════════════════════════════════════════════╬════════════╬═══════════╬══════════╬══════════════╣
║  3 ║ Calculator::prepare                          ║          2 ║         4 ║  53.29ms ║      26.64ms ║
╠════╬══════════════════════════════════════════════╬════════════╬═══════════╬══════════╬══════════════╣
║  2 ║ NeighborsList                                ║         40 ║         3 ║  11.93ms ║     298.34µs ║
╠════╬══════════════════════════════════════════════╬════════════╬═══════════╬══════════╬══════════════╣
║  6 ║ SplinedRadialIntegral::compute               ║      35307 ║         5 ║  17.01ms ║ 481.00ns ⚠️  ║
╠════╬══════════════════════════════════════════════╬════════════╬═══════════╬══════════╬══════════════╣
║  7 ║ SphericalHarmonics::compute                  ║      35307 ║         5 ║  20.30ms ║ 574.00ns ⚠️  ║
╚════╩══════════════════════════════════════════════╩════════════╩═══════════╩══════════╩══════════════╝

You can also save this data as json for future usage

print(profiler.as_json())

{"timings":{"rascaline::calculators::soap::radial_integral::gto::GtoRadialIntegral::compute":{"id":0,"elapsed":"50.157631ms","called":3846},"rascaline::calculators::soap::radial_integral::spline::SoapRadialIntegralSpline::with_accuracy":{"id":1,"elapsed":"53.028376ms","called":6},"rascaline::systems::neighbors::NeighborsList":{"id":2,"elapsed":"11.93377ms","called":40},"rascaline::calculator::Calculator::prepare":{"id":3,"elapsed":"53.286337ms","called":2},"rascaline::calculators::soap::power_spectrum::SoapPowerSpectrum::compute":{"id":4,"elapsed":"420.089266ms","called":1},"rascaline::calculators::soap::spherical_expansion::SphericalExpansion::compute":{"id":5,"elapsed":"197.608947ms","called":1},"rascaline::calculators::soap::radial_integral::spline::SplinedRadialIntegral::compute":{"id":6,"elapsed":"17.007453ms","called":35307},"rascaline::math::spherical_harmonics::SphericalHarmonics::compute":{"id":7,"elapsed":"20.300349ms","called":35307}},"calls":[{"caller":0,"callee":1,"count":3846},{"caller":2,"callee":3,"count":40},{"caller":3,"callee":4,"count":1},{"caller":1,"callee":5,"count":1},{"caller":6,"callee":5,"count":1},{"caller":7,"callee":5,"count":1},{"caller":5,"callee":4,"count":1}]}

use metatensor::{TensorMap, Labels};
use rascaline::{Calculator, System, CalculationOptions};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let path = std::env::args().nth(1).expect("expected a command line argument");

    // enable collection of profiling data
    time_graph::enable_data_collection(true);
    // clear any existing collected data
    time_graph::clear_collected_data();

    // run the calculation
    let _descriptor = compute_soap(&path)?;

    // get the call graph and display it
    let graph = time_graph::get_full_graph();
    // (this requires the "table" feature for the time_graph crate)
    println!("{}", graph.as_short_table());

    // also available for saving profiling data to the disk & future analysis
    // (this requires the "json" feature for the time_graph crate)
    println!("{}", graph.as_json());

    Ok(())
}


/// Compute SOAP power spectrum, this is the same code as the 'compute-soap'
/// example
fn compute_soap(path: &str) -> Result<TensorMap, Box<dyn std::error::Error>> {
    let systems = rascaline::systems::read_from_file(path)?;
    let mut systems = systems.into_iter()
        .map(|s| Box::new(s) as Box<dyn System>)
        .collect::<Vec<_>>();

    let parameters = r#"{
        "cutoff": 5.0,
        "max_radial": 6,
        "max_angular": 4,
        "atomic_gaussian_width": 0.3,
        "center_atom_weight": 1.0,
        "radial_basis": {
            "Gto": {}
        },
        "cutoff_function": {
            "ShiftedCosine": {"width": 0.5}
        }
    }"#;

    let descriptor = time_graph::spanned!("Full calculation", {
        let mut calculator = Calculator::new("soap_power_spectrum", parameters.to_owned())?;

        let options = CalculationOptions {
            gradients: &["positions"],
            ..Default::default()
        };
        calculator.compute(&mut systems, options)?
    });

    let keys_to_move = Labels::empty(vec!["center_type"]);
    let descriptor = descriptor.keys_to_samples(&keys_to_move, /* sort_samples */ true)?;

    let keys_to_move = Labels::empty(vec!["neighbor_1_type", "neighbor_2_type"]);
    let descriptor = descriptor.keys_to_properties(&keys_to_move, /* sort_samples */ true)?;

    Ok(descriptor)
}

#include <iostream>
#include <rascaline.hpp>

/// Compute SOAP power spectrum, this is the same code as the 'compute-soap'
/// example
static metatensor::TensorMap compute_soap(const std::string& path);

int main(int argc, char* argv[]) {
    if (argc < 2) {
        std::cout << "error: expected a command line argument" << std::endl;
        return 1;
    }

    // enable collection of profiling data
    rascaline::Profiler::enable(true);
    // clear any existing collected data
    rascaline::Profiler::clear();

    auto descriptor = compute_soap(argv[1]);

    // Get the profiling data as a table to display it directly
    std::cout << rascaline::Profiler::get("short_table") << std::endl;
    // Or save this data as json for future usage
    std::cout << rascaline::Profiler::get("json") << std::endl;

    return 0;
}


metatensor::TensorMap compute_soap(const std::string& path) {
    auto systems = rascaline::BasicSystems(path);

    const char* parameters = R"({
        "cutoff": 5.0,
        "max_radial": 6,
        "max_angular": 4,
        "atomic_gaussian_width": 0.3,
        "center_atom_weight": 1.0,
        "gradients": false,
        "radial_basis": {
            "Gto": {}
        },
        "cutoff_function": {
            "ShiftedCosine": {"width": 0.5}
        }
    })";

    auto calculator = rascaline::Calculator("soap_power_spectrum", parameters);
    auto descriptor = calculator.compute(systems);

    descriptor.keys_to_samples("center_type");
    descriptor.keys_to_properties(std::vector<std::string>{"neighbor_1_type", "neighbor_2_type"});

    return descriptor;
}

#include <stdbool.h>
#include <stdio.h>

#include <metatensor.h>
#include <rascaline.h>

/// Compute SOAP power spectrum, this is the same code as the 'compute-soap'
/// example
static mts_tensormap_t* compute_soap(const char* path);

int main(int argc, char* argv[]) {
    rascal_status_t status = RASCAL_SUCCESS;
    char* buffer = NULL;
    size_t buffer_size = 8192;
    bool got_error = true;

    if (argc < 2) {
        printf("error: expected a command line argument");
        goto cleanup;
    }

    // enable collection of profiling data
    status = rascal_profiling_enable(true);
    if (status != RASCAL_SUCCESS) {
        printf("Error: %s\n", rascal_last_error());
        goto cleanup;
    }

    // clear any existing collected data
    status = rascal_profiling_clear();
    if (status != RASCAL_SUCCESS) {
        printf("Error: %s\n", rascal_last_error());
        goto cleanup;
    }

    mts_tensormap_t* descriptor = compute_soap(argv[1]);
    if (descriptor == NULL) {
        goto cleanup;
    }

    buffer = calloc(buffer_size, sizeof(char));
    if (buffer == NULL) {
        printf("Error: failed to allocate memory\n");
        goto cleanup;
    }

    // Get the profiling data as a table to display it directly
    status = rascal_profiling_get("short_table", buffer, buffer_size);
    if (status != RASCAL_SUCCESS) {
        printf("Error: %s\n", rascal_last_error());
        goto cleanup;
    }
    printf("%s\n", buffer);

    // Or save this data as json for future usage
    status = rascal_profiling_get("json", buffer, buffer_size);
    if (status != RASCAL_SUCCESS) {
        printf("Error: %s\n", rascal_last_error());
        goto cleanup;
    }
    printf("%s\n", buffer);

    got_error = false;
cleanup:
    free(buffer);
    mts_tensormap_free(descriptor);

    if (got_error) {
        return 1;
    } else {
        return 0;
    }
}

static mts_tensormap_t* move_keys_to_samples(mts_tensormap_t* descriptor, const char* keys_to_move[], size_t keys_to_move_len);
static mts_tensormap_t* move_keys_to_properties(mts_tensormap_t* descriptor, const char* keys_to_move[], size_t keys_to_move_len);

// this is the same function as in the compute-soap.c example
mts_tensormap_t* compute_soap(const char* path) {
    int status = RASCAL_SUCCESS;
    rascal_calculator_t* calculator = NULL;
    rascal_system_t* systems = NULL;
    uintptr_t n_systems = 0;
    const double* values = NULL;
    const uintptr_t* shape = NULL;
    uintptr_t shape_count = 0;
    bool got_error = true;
    const char* keys_to_samples[] = {"center_type"};
    const char* keys_to_properties[] = {"neighbor_1_type", "neighbor_2_type"};

    // use the default set of options, computing all samples and all features
    rascal_calculation_options_t options = {0};
    const char* gradients_list[] = {"positions"};
    options.gradients = gradients_list;
    options.gradients_count = 1;

    mts_tensormap_t* descriptor = NULL;
    const mts_block_t* block = NULL;
    mts_array_t data = {0};
    mts_labels_t keys_to_move = {0};

    const char* parameters = "{\n"
        "\"cutoff\": 5.0,\n"
        "\"max_radial\": 6,\n"
        "\"max_angular\": 4,\n"
        "\"atomic_gaussian_width\": 0.3,\n"
        "\"center_atom_weight\": 1.0,\n"
        "\"gradients\": false,\n"
        "\"radial_basis\": {\n"
        "    \"Gto\": {}\n"
        "},\n"
        "\"cutoff_function\": {\n"
        "    \"ShiftedCosine\": {\"width\": 0.5}\n"
        "}\n"
    "}";


    status = rascal_basic_systems_read(path, &systems, &n_systems);
    if (status != RASCAL_SUCCESS) {
        printf("Error: %s\n", rascal_last_error());
        goto cleanup;
    }

    calculator = rascal_calculator("soap_power_spectrum", parameters);
    if (calculator == NULL) {
        printf("Error: %s\n", rascal_last_error());
        goto cleanup;
    }

    status = rascal_calculator_compute(
        calculator, &descriptor, systems, n_systems, options
    );
    if (status != RASCAL_SUCCESS) {
        printf("Error: %s\n", rascal_last_error());
        goto cleanup;
    }

    descriptor = move_keys_to_samples(descriptor, keys_to_samples, 1);
    if (descriptor == NULL) {
        printf("Error: %s\n", mts_last_error());
        goto cleanup;
    }

    descriptor = move_keys_to_properties(descriptor, keys_to_properties, 2);
    if (descriptor == NULL) {
        printf("Error: %s\n", mts_last_error());
        goto cleanup;
    }

cleanup:
    rascal_calculator_free(calculator);
    rascal_basic_systems_free(systems, n_systems);

    return descriptor;
}


mts_tensormap_t* move_keys_to_samples(mts_tensormap_t* descriptor, const char* keys_to_move[], size_t keys_to_move_len) {
    mts_labels_t keys = {0};
    mts_tensormap_t* moved_descriptor = NULL;

    keys.names = keys_to_move;
    keys.size = keys_to_move_len;
    keys.values = NULL;
    keys.count = 0;

    moved_descriptor = mts_tensormap_keys_to_samples(descriptor, keys, true);
    mts_tensormap_free(descriptor);

    return moved_descriptor;
}


mts_tensormap_t* move_keys_to_properties(mts_tensormap_t* descriptor, const char* keys_to_move[], size_t keys_to_move_len) {
    mts_labels_t keys = {0};
    mts_tensormap_t* moved_descriptor = NULL;

    keys.names = keys_to_move;
    keys.size = keys_to_move_len;
    keys.values = NULL;
    keys.count = 0;

    moved_descriptor = mts_tensormap_keys_to_properties(descriptor, keys, true);
    mts_tensormap_free(descriptor);

    return moved_descriptor;
}