mirror of
https://github.com/gcorso/DiffDock.git
synced 2026-06-04 18:04:23 +08:00
177 lines
3.3 KiB
Python
177 lines
3.3 KiB
Python
import os
|
|
|
|
import numpy as np
|
|
|
|
|
|
def format(fname: str) -> str:
|
|
"""
|
|
Extract format extension from file name.
|
|
|
|
Parameters
|
|
----------
|
|
fname : str
|
|
File name
|
|
|
|
Returns
|
|
-------
|
|
str
|
|
File extension
|
|
|
|
Notes
|
|
-----
|
|
The file extension is returned without the `.` character, i.e. for the file
|
|
`path/filename.ext` the string `ext` is returned.
|
|
|
|
If a file is compressed, the `.gz` extension is ignored.
|
|
"""
|
|
name, ext = os.path.splitext(fname)
|
|
|
|
if ext == ".gz":
|
|
_, ext = os.path.splitext(name)
|
|
|
|
return ext[1:] # Remove "."
|
|
|
|
|
|
def molformat(fname: str) -> str:
|
|
"""
|
|
Extract an OpenBabel-friendly format from file name.
|
|
|
|
Parameters
|
|
----------
|
|
fname : str
|
|
File name
|
|
|
|
Returns
|
|
-------
|
|
str
|
|
File extension in an OpenBabel-friendly format
|
|
|
|
Notes
|
|
-----
|
|
File types in OpenBabel do not always correspond to the file extension. This
|
|
function converts the file extension to an OpenBabel file type.
|
|
|
|
The following table shows the different conversions performed by this function:
|
|
|
|
========= =========
|
|
Extension File Type
|
|
--------- ---------
|
|
xyz XYZ
|
|
========= =========
|
|
"""
|
|
|
|
ext = format(fname)
|
|
|
|
if ext == "xyz":
|
|
# xyz files in OpenBabel are called XYZ
|
|
ext = "XYZ"
|
|
|
|
return ext
|
|
|
|
|
|
def deg_to_rad(angle: float) -> float:
|
|
"""
|
|
Convert angle in degrees to angle in radians.
|
|
|
|
Parameters
|
|
----------
|
|
angle : float
|
|
Angle (in degrees)
|
|
|
|
Returns
|
|
-------
|
|
float
|
|
Angle (in radians)
|
|
"""
|
|
|
|
return angle * np.pi / 180.0
|
|
|
|
|
|
def rotate(
|
|
v: np.ndarray, angle: float, axis: np.ndarray, units: str = "rad"
|
|
) -> np.ndarray:
|
|
"""
|
|
Rotate vector.
|
|
|
|
Parameters
|
|
----------
|
|
v: numpy.array
|
|
3D vector to be rotated
|
|
angle : float
|
|
Angle of rotation (in `units`)
|
|
axis : numpy.array
|
|
3D axis of rotation
|
|
units: {"rad", "deg"}
|
|
Units of `angle` (in radians `rad` or degrees `deg`)
|
|
|
|
Returns
|
|
-------
|
|
numpy.array
|
|
Rotated vector
|
|
|
|
Raises
|
|
------
|
|
AssertionError
|
|
If the axis of rotation is not a 3D vector
|
|
ValueError
|
|
If `units` is not `rad` or `deg`
|
|
"""
|
|
|
|
assert len(axis) == 3
|
|
|
|
# Ensure rotation axis is normalised
|
|
axis = axis / np.linalg.norm(axis)
|
|
|
|
if units.lower() == "rad":
|
|
pass
|
|
elif units.lower() == "deg":
|
|
angle = deg_to_rad(angle)
|
|
else:
|
|
raise ValueError(
|
|
f"Units {units} for angle is not supported. Use 'deg' or 'rad' instead."
|
|
)
|
|
|
|
t1 = np.outer(axis, np.inner(axis, v)).T
|
|
t2 = np.cos(angle) * np.cross(np.cross(axis, v), axis)
|
|
t3 = np.sin(angle) * np.cross(axis, v)
|
|
|
|
return t1 + t2 + t3
|
|
|
|
|
|
def center_of_geometry(coordinates: np.ndarray) -> np.ndarray:
|
|
"""
|
|
Center of geometry.
|
|
|
|
Parameters
|
|
----------
|
|
coordinates: np.ndarray
|
|
Coordinates
|
|
|
|
Returns
|
|
-------
|
|
np.ndarray
|
|
Center of geometry
|
|
"""
|
|
|
|
assert coordinates.shape[1] == 3
|
|
|
|
return np.mean(coordinates, axis=0)
|
|
|
|
|
|
def center(coordinates: np.ndarray) -> np.ndarray:
|
|
"""
|
|
Center coordinates.
|
|
|
|
Parameters
|
|
----------
|
|
coordinates: np.ndarray
|
|
Coordinates
|
|
|
|
Returns
|
|
-------
|
|
np.ndarray
|
|
Centred coordinates
|
|
"""
|
|
|
|
return coordinates - center_of_geometry(coordinates)
|