Skip to content

Commit 0850575

Browse files
lohedgeslohedges
authored
Merge pull request #465 from OpenBioSim/feature_set_coordinates
Add support for setting coordinates of all atoms from a NumPy array
2 parents 2cc1724 + 099db38 commit 0850575

File tree

4 files changed

+274
-2
lines changed

4 files changed

+274
-2
lines changed

‎python/BioSimSpace/Sandpit/Exscientia/_SireWrappers/_system.py‎

Lines changed: 112 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -380,7 +380,7 @@ def charge(self, property_map={}, is_lambda1=False):
380380
own naming scheme, e.g. { "charge" : "my-charge" }
381381
382382
is_lambda1 : bool
383-
Whether to use the charge at lambda = 1 if the molecule is merged.
383+
Whether to use the charge at lambda = 1 for perturbable molecules.
384384
385385
Returns
386386
-------
@@ -1210,6 +1210,117 @@ def nMLMolecules(self):
12101210
"""
12111211
return len(self.getMLMolecules())
12121212

1213+
def getCoordinates(self, is_lambda1=False, property_map={}):
1214+
"""
1215+
Return the coordinates of all atoms in the system as a NumPy array.
1216+
Coordinates are returned in Angstroms.
1217+
1218+
Parameters
1219+
----------
1220+
1221+
is_lambda1 : bool
1222+
Whether to use the coordinates at lambda = 1 for perturbable molecules.
1223+
1224+
property_map : dict
1225+
A dictionary that maps system "properties" to their user defined
1226+
values. This allows the user to refer to properties with their
1227+
own naming scheme, e.g. { "charge" : "my-charge" }
1228+
1229+
Returns
1230+
-------
1231+
1232+
coordinates : numpy.ndarray
1233+
The coordinates of all atoms in the system in Angstroms.
1234+
"""
1235+
1236+
if not isinstance(is_lambda1, bool):
1237+
raise TypeError("'is_lambda1' must be of type 'bool'")
1238+
1239+
if not isinstance(property_map, dict):
1240+
raise TypeError("'property_map' must be of type 'dict'")
1241+
1242+
import sire as _sr
1243+
1244+
# Convert to a new Sire system.
1245+
mols = _sr.system.System(self._sire_object)
1246+
1247+
# Link to the correct end-state if required.
1248+
if self.nPerturbableMolecules() > 0:
1249+
if is_lambda1:
1250+
mols = _sr.morph.link_to_perturbed(mols, map=property_map)
1251+
else:
1252+
mols = _sr.morph.link_to_reference(mols, map=property_map)
1253+
1254+
# Try to get the coordinates array.
1255+
try:
1256+
coords = _sr.io.get_coords_array(mols, map=property_map)
1257+
except Exception as e:
1258+
msg = "Failed to extract coordinates from system!"
1259+
if _isVerbose():
1260+
raise RuntimeError(msg) from e
1261+
else:
1262+
raise RuntimeError(msg) from None
1263+
1264+
return coords
1265+
1266+
def setCoordinates(
1267+
self,
1268+
coordinates,
1269+
is_lambda1=False,
1270+
property_map={},
1271+
):
1272+
"""
1273+
Set the coordinates of all atoms in the system from a NumPy array.
1274+
Coordinates are expected to be in Angstroms.
1275+
1276+
Parameters
1277+
----------
1278+
1279+
coordinates : numpy.ndarray
1280+
The coordinates of all atoms in the system in Angstroms.
1281+
1282+
is_lambda1 : bool
1283+
Whether to set the coordinates at lambda = 1 for perturbable molecules.
1284+
1285+
property_map : dict
1286+
A dictionary that maps system "properties" to their user defined
1287+
values. This allows the user to refer to properties with their
1288+
own naming scheme, e.g. { "charge" : "my-charge" }
1289+
"""
1290+
1291+
import numpy as _np
1292+
1293+
# Validate input.
1294+
if not isinstance(coordinates, _np.ndarray):
1295+
raise TypeError("'coordinates' must be of type 'numpy.ndarray'")
1296+
if coordinates.ndim != 2 or coordinates.shape[1] != 3:
1297+
raise ValueError("'coordinates' must be a 2D array with shape (n_atoms, 3)")
1298+
if coordinates.shape[0] != self.nAtoms():
1299+
raise ValueError(
1300+
"'coordinates' must have the same number of atoms as the system"
1301+
)
1302+
1303+
if not isinstance(is_lambda1, bool):
1304+
raise TypeError("'is_lambda1' must be of type 'bool'")
1305+
1306+
if not isinstance(property_map, dict):
1307+
raise TypeError("'property_map' must be of type 'dict'")
1308+
1309+
# Set the coordinates.
1310+
try:
1311+
self._sire_object = _SireIO.setCoordinates(
1312+
self._sire_object,
1313+
coordinates.tolist(),
1314+
is_lambda1,
1315+
map=property_map,
1316+
)
1317+
except Exception as e:
1318+
msg = "Failed to set coordinates in system!"
1319+
if _isVerbose():
1320+
raise RuntimeError(msg) from e
1321+
else:
1322+
raise RuntimeError(msg) from None
1323+
12131324
def rotateBoxVectors(
12141325
self,
12151326
origin=_Coordinate(

‎python/BioSimSpace/_SireWrappers/_system.py‎

Lines changed: 112 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -380,7 +380,7 @@ def charge(self, property_map={}, is_lambda1=False):
380380
own naming scheme, e.g. { "charge" : "my-charge" }
381381
382382
is_lambda1 : bool
383-
Whether to use the charge at lambda = 1 if the molecule is merged.
383+
Whether to use the charge at lambda = 1 for perturbable molecules.
384384
385385
Returns
386386
-------
@@ -1158,6 +1158,117 @@ def nPerturbableMolecules(self):
11581158
"""
11591159
return len(self.getPerturbableMolecules())
11601160

1161+
def getCoordinates(self, is_lambda1=False, property_map={}):
1162+
"""
1163+
Return the coordinates of all atoms in the system as a NumPy array.
1164+
Coordinates are returned in Angstroms.
1165+
1166+
Parameters
1167+
----------
1168+
1169+
is_lambda1 : bool
1170+
Whether to use the coordinates at lambda = 1 for perturbable molecules.
1171+
1172+
property_map : dict
1173+
A dictionary that maps system "properties" to their user defined
1174+
values. This allows the user to refer to properties with their
1175+
own naming scheme, e.g. { "charge" : "my-charge" }
1176+
1177+
Returns
1178+
-------
1179+
1180+
coordinates : numpy.ndarray
1181+
The coordinates of all atoms in the system in Angstroms.
1182+
"""
1183+
1184+
if not isinstance(is_lambda1, bool):
1185+
raise TypeError("'is_lambda1' must be of type 'bool'")
1186+
1187+
if not isinstance(property_map, dict):
1188+
raise TypeError("'property_map' must be of type 'dict'")
1189+
1190+
import sire as _sr
1191+
1192+
# Convert to a new Sire system.
1193+
mols = _sr.system.System(self._sire_object)
1194+
1195+
# Link to the correct end-state if required.
1196+
if self.nPerturbableMolecules() > 0:
1197+
if is_lambda1:
1198+
mols = _sr.morph.link_to_perturbed(mols, map=property_map)
1199+
else:
1200+
mols = _sr.morph.link_to_reference(mols, map=property_map)
1201+
1202+
# Try to get the coordinates array.
1203+
try:
1204+
coords = _sr.io.get_coords_array(mols, map=property_map)
1205+
except Exception as e:
1206+
msg = "Failed to extract coordinates from system!"
1207+
if _isVerbose():
1208+
raise RuntimeError(msg) from e
1209+
else:
1210+
raise RuntimeError(msg) from None
1211+
1212+
return coords
1213+
1214+
def setCoordinates(
1215+
self,
1216+
coordinates,
1217+
is_lambda1=False,
1218+
property_map={},
1219+
):
1220+
"""
1221+
Set the coordinates of all atoms in the system from a NumPy array.
1222+
Coordinates are expected to be in Angstroms.
1223+
1224+
Parameters
1225+
----------
1226+
1227+
coordinates : numpy.ndarray
1228+
The coordinates of all atoms in the system in Angstroms.
1229+
1230+
is_lambda1 : bool
1231+
Whether to set the coordinates at lambda = 1 for perturbable molecules.
1232+
1233+
property_map : dict
1234+
A dictionary that maps system "properties" to their user defined
1235+
values. This allows the user to refer to properties with their
1236+
own naming scheme, e.g. { "charge" : "my-charge" }
1237+
"""
1238+
1239+
import numpy as _np
1240+
1241+
# Validate input.
1242+
if not isinstance(coordinates, _np.ndarray):
1243+
raise TypeError("'coordinates' must be of type 'numpy.ndarray'")
1244+
if coordinates.ndim != 2 or coordinates.shape[1] != 3:
1245+
raise ValueError("'coordinates' must be a 2D array with shape (n_atoms, 3)")
1246+
if coordinates.shape[0] != self.nAtoms():
1247+
raise ValueError(
1248+
"'coordinates' must have the same number of atoms as the system"
1249+
)
1250+
1251+
if not isinstance(is_lambda1, bool):
1252+
raise TypeError("'is_lambda1' must be of type 'bool'")
1253+
1254+
if not isinstance(property_map, dict):
1255+
raise TypeError("'property_map' must be of type 'dict'")
1256+
1257+
# Set the coordinates.
1258+
try:
1259+
self._sire_object = _SireIO.setCoordinates(
1260+
self._sire_object,
1261+
coordinates.tolist(),
1262+
is_lambda1,
1263+
map=property_map,
1264+
)
1265+
except Exception as e:
1266+
msg = "Failed to set coordinates in system!"
1267+
if _isVerbose():
1268+
raise RuntimeError(msg) from e
1269+
else:
1270+
raise RuntimeError(msg) from None
1271+
11611272
def rotateBoxVectors(
11621273
self,
11631274
origin=_Coordinate(

‎tests/Sandpit/Exscientia/_SireWrappers/test_system.py‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import math
2+
import numpy as np
23
import pytest
34

45
from sire.legacy.Vol import TriclinicBox
@@ -531,3 +532,27 @@ def test_renumber(system):
531532

532533
# Make sure that no original numbers are present in the renumbered set.
533534
assert original_numbers.isdisjoint(renumbered_numbers)
535+
536+
537+
@pytest.mark.parametrize("fixture", ["system", "perturbable_system"])
538+
@pytest.mark.parametrize("is_lambda1", [True, False])
539+
def test_set_coordinates(fixture, is_lambda1, request):
540+
# Get the fixture system.
541+
mols = request.getfixturevalue(fixture).copy()
542+
543+
# Store the existing coordinates as a NumPy array.
544+
coords = mols.getCoordinates(is_lambda1=is_lambda1)
545+
546+
# Modify the system to multiply all coordinates by 2.
547+
mols.setCoordinates(coords * 2.0, is_lambda1=is_lambda1)
548+
549+
# Get the new coordinates as a NumPy array.
550+
new_coords = mols.getCoordinates(is_lambda1=is_lambda1)
551+
552+
# Divide the new coordinates by the old coordinates.
553+
ratio = new_coords / coords
554+
555+
# Make sure the new coordinates are as expected.
556+
assert (
557+
np.sum(np.round(ratio)) == 6.0 * mols.nAtoms()
558+
), "Coordinates were not set correctly."

‎tests/_SireWrappers/test_system.py‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import math
2+
import numpy as np
23
import pytest
34

45
from sire.legacy.Vol import TriclinicBox
@@ -521,3 +522,27 @@ def test_renumber(system):
521522

522523
# Make sure that no original numbers are present in the renumbered set.
523524
assert original_numbers.isdisjoint(renumbered_numbers)
525+
526+
527+
@pytest.mark.parametrize("fixture", ["system", "perturbable_system"])
528+
@pytest.mark.parametrize("is_lambda1", [True, False])
529+
def test_set_coordinates(fixture, is_lambda1, request):
530+
# Get the fixture system.
531+
mols = request.getfixturevalue(fixture).copy()
532+
533+
# Store the existing coordinates as a NumPy array.
534+
coords = mols.getCoordinates(is_lambda1=is_lambda1)
535+
536+
# Modify the system to multiply all coordinates by 2.
537+
mols.setCoordinates(coords * 2.0, is_lambda1=is_lambda1)
538+
539+
# Get the new coordinates as a NumPy array.
540+
new_coords = mols.getCoordinates(is_lambda1=is_lambda1)
541+
542+
# Divide the new coordinates by the old coordinates.
543+
ratio = new_coords / coords
544+
545+
# Make sure the new coordinates are as expected.
546+
assert (
547+
np.sum(np.round(ratio)) == 6.0 * mols.nAtoms()
548+
), "Coordinates were not set correctly."

0 commit comments

Comments
 (0)