Local flask server for espaloma charge

Last week I wanted some reasonable partial charges to help with an atom typing experiment. The Chodera lab's Espaloma Charge model is super fast and apparently high quality (an alternative to use might be kallisto). Trying to install Espaloma immediately ruined my conda env thanks to conflicts with DGL/pytorch versions, so I decided it's best to put it in a separate, cordoned off, environment. So now I've got two environments - one with all my working libraries, and a separate one just for espaloma. How do I get molecule structure information out of the former, and into the latter, and then return the partial charges? I don't want to write the SMILES/SDF to file, stop, instantiate the other environment, parse the file, write out the charges, move back, and then load the charges. So many fiddly steps! And I want it to be interactive.

The answer was to set up a little local server, using Flask and running Espaloma, to which I could POST a SMILES code or SDF block, allow the server to run the espaloma charging function, and return the partial charges as a JSON.

Setting up this espaloma/flask server looks like this:

from flask import Flask, request, jsonify
import numpy as np
from rdkit import Chem

from espaloma_charge import charge

app = Flask(__name__)

@app.route('/process', methods=['POST'])
def process():
    data = request.json
    input_string = data['input']
    m = Chem.MolFromMolBlock(input_string, removeHs=False) # do NOT remove Hs ;)
    result = charge(m)
    return jsonify(result.tolist())

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)
  

Run it in the shell, with the espaloma environment activated, and let it sit there to wait patiently for POST'd data. Now, in the working environment (for me a notebook), run:

import requests
import json
response = requests.post('http://localhost:5000/process', json={'input': Chem.MolToMolBlock(m)})
partial_charges = np.array(response.json())
  

et voila:

array([-0.19000548, -0.19000548,  0.63194072,  0.63194072,  0.45600772,
   -0.49415129,  0.34710342, -0.20972057,  0.49532843, -1.00420523,
   -0.6615752 , -0.6615752 , -0.34363997, -0.28776151,  1.65871096,
   -0.17839202])