a5ee2e0ea284a9b024721536abb94b7998a9e8ee,qcengine/programs/torchani.py,,run_torchani,#Any#,8

Before Change

    ret_data["success"] = False

    // Build model and species array
    if ret_data["model"]["method"] == "ANI1":

        // Build model
        builtin = torchani.neurochem.Builtins()
        model = torch.nn.Sequential(builtin.aev_computer, builtin.models, builtin.energy_shifter)

        // Build species
        unknown_sym = set(ret_data["molecule"]["symbols"]) - {"H", "C", "N", "O"}

After Change

    // Build model
    model = get_model(ret_data["model"]["method"])
    if model is False:
        ret_data["error_message"] = "run_torchani only accepts the ANI1 method."
        return ret_data

    // Build species
    species = "".join(ret_data["molecule"]["symbols"])
    unknown_sym = set(species) - {"H", "C", "N", "O"}
    if unknown_sym:
        ret_data["error_message"] = "The "{}" model does not support symbols: {}.".format(
            ret_data["model"]["method"], unknown_sym)
        return ret_data

    species = builtin.consts.species_to_tensor(species).to(device).unsqueeze(0)

    // Build coord array
    geom_array = np.array(ret_data["molecule"]["geometry"]).reshape(1, -1, 3) * units.bohr_to_angstrom
    coordinates = torch.tensor(geom_array.tolist(), requires_grad=True, device=device)

    _, energy = model((species, coordinates))
    ret_data["properties"] = {"return_energy": energy.item()}

    if ret_data["driver"] == "energy":
        ret_data["return_result"] = ret_data["properties"]["return_energy"]
    elif ret_data["driver"] == "gradient":
        derivative = torch.autograd.grad(energy.sum(), coordinates)[0].squeeze()
        ret_data["return_result"] = np.asarray(derivative / units.bohr_to_angstrom).ravel().tolist()
    else:
        ret_data["error_message"] = "run_torchani did not understand driver method "{}".".format(ret_data["driver"])
        return ret_data

Instances