simple_edm_to_coords

from rdkit import Chem
import numpy as np
from rdkit.Chem import AllChem
from rdkit.Geometry import Point3D

from scipy.spatial.distance import squareform, pdist

import copy

import py3Dmol

def to_M_matrix(D):
    D1j = np.tile( D[0,:], len(D)).reshape(D.shape)
    Di1 = D1j.T
    return 0.5 * (-D + D1j + Di1)

def get_coords(D):
    M = to_M_matrix(D**2)
    S, U = np.linalg.eigh(M)
    
    coord = np.matmul(U, np.diag(np.sqrt(np.abs(S))))
    return coord[:, -3:]

def setPositions(mol_in, coords):
    mol = copy.copy(mol_in)
    conformer = mol.GetConformer(0)
    for i in range(mol.GetNumAtoms()):
        x,y,z=coords[i]
        conformer.SetAtomPosition(i, Point3D(x,y,z))
    return mol


def viewProt(mol, view=None, color='orange'):
    if view is None:
        view = py3Dmol.view()
    view.addModel(Chem.MolToMolBlock(mol), "sdf")

    model = view.getModel()
    model.setStyle({}, {"stick": {"colorscheme": f"{color}Carbon"}})
    view.zoomTo({})
    return view


lig = Chem.MolFromSmiles('Cc1c(c2cc(ccc2n1C(=O)c3ccc(cc3)Cl)OC)CC(=O)O')


ligH = Chem.AddHs(lig)
AllChem.EmbedMolecule(ligH)
lig = Chem.RemoveHs(ligH)

lpos = lig.GetConformer(0).GetPositions()

ldmat = squareform(pdist(lpos))

X =get_coords(ldmat)
newmol = setPositions(lig, X)

viewProt(newmol)