forked from enviPath/enviPy
84 lines
2.1 KiB
Python
84 lines
2.1 KiB
Python
import json
|
|
|
|
import requests
|
|
|
|
|
|
class AMBITResult:
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
self.smiles = kwargs['smiles']
|
|
self.tps = []
|
|
for bt in kwargs['products']:
|
|
if len(bt['products']):
|
|
self.tps.append(bt)
|
|
|
|
self.probs = None
|
|
|
|
def __str__(self):
|
|
x = self.smiles + "\n"
|
|
total_bts = len(self.tps)
|
|
|
|
for i, tp in enumerate(self.tps):
|
|
prob = ""
|
|
if self.probs:
|
|
prob = f" (p={self.probs[tp['id']]})"
|
|
|
|
if i == total_bts - 1:
|
|
x += f"\t└── {tp['name']}{prob}\n"
|
|
else:
|
|
x += f"\t├── {tp['name']}{prob}\n"
|
|
|
|
total_products = len(tp['products'])
|
|
for j, p in enumerate(tp['products']):
|
|
if j == total_products - 1:
|
|
if i == total_bts - 1:
|
|
x += f"\t\t└── {p}"
|
|
else:
|
|
x += f"\t│\t└── {p}\n"
|
|
else:
|
|
if i == total_bts - 1:
|
|
x += f"\t\t├── {p}\n"
|
|
else:
|
|
x += f"\t│\t├── {p}\n"
|
|
return x
|
|
|
|
def set_probs(self, probs):
|
|
self.probs = probs
|
|
|
|
|
|
class AMBIT:
|
|
|
|
def __init__(self, host, rules=None):
|
|
self.host = host
|
|
self.rules = rules
|
|
self.ambit_params = {
|
|
'singlePos': True,
|
|
'split': False,
|
|
}
|
|
|
|
def batch_apply(self, smiles: list):
|
|
payload = {
|
|
'smiles': smiles,
|
|
'rules': self.rules,
|
|
}
|
|
payload.update(**self.ambit_params)
|
|
|
|
res = self._execute(payload)
|
|
|
|
tps = list()
|
|
for r in res['result']:
|
|
ar = AMBITResult(**r)
|
|
if len(ar.tps):
|
|
tps.append(ar)
|
|
else:
|
|
tps.append(None)
|
|
return tps
|
|
|
|
def apply(self, smiles: str):
|
|
return self.batch_apply([smiles])[0]
|
|
|
|
def _execute(self, payload):
|
|
res = requests.post(self.host + '/ambit', data=json.dumps(payload))
|
|
res.raise_for_status()
|
|
return res.json()
|