Skip to content

openmm

deepdrivewe.simulation.openmm

OpenMM simulation module.

OpenMMReporter

Bases: ABC

Reporter interface for OpenMM simulations.

Source code in deepdrivewe/simulation/openmm.py 
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
class OpenMMReporter(ABC):
    """Reporter interface for OpenMM simulations."""

    def __init__(
        self,
        report_interval: int,
        openmm_selection: Sequence[str] = ('CA',),
    ) -> None:
        """Initialize the reporter.

        Parameters
        ----------
        report_interval : int
            The interval at which to write frames.
        openmm_selection : Sequence[str]
            The OpenMM selection strings for the atoms to use
            when reporting positions (default is ('CA',)).
        """
        self.report_interval = report_interval
        self.openmm_selection = openmm_selection

    def describeNextReport(  # noqa: N802
        self,
        simulation: app.Simulation,
    ) -> tuple[int, bool, bool, bool, bool, bool | None]:
        """Get information about the next report this object will generate.

        Parameters
        ----------
        simulation : Simulation
            The Simulation to generate a report for

        Returns
        -------
        tuple
            A six element tuple. The first element is the number of steps
            until the next report. The next four elements specify whether
            that report will require positions, velocities, forces, and
            energies respectively.  The final element specifies whether
            positions should be wrapped to lie in a single periodic box.
        """
        step_progress = simulation.currentStep % self.report_interval
        steps = self.report_interval - step_progress
        return (steps, True, False, False, False, None)

    def get_positions(
        self,
        simulation: app.Simulation,
        state: openmm.State,
    ) -> np.ndarray:
        """Get the atomic positions from the simulation.

        Parameters
        ----------
        simulation : Simulation
            The Simulation to generate a report for.
        state : State
            The current state of the simulation.

        Returns
        -------
        np.ndarray
            The atomic positions from the simulation.
        """
        # Get the atom indices for the selection
        atom_indices = [
            a.index
            for a in simulation.topology.atoms()
            if a.name in self.openmm_selection
        ]

        # Get the atomic coordinates of the selection
        positions = state.getPositions(asNumpy=True)
        positions = positions[atom_indices].astype(np.float32)

        # Convert positions from nanometers to angstroms
        positions *= 10.0

        return positions

    @abstractmethod
    def report(self, simulation: app.Simulation, state: openmm.State) -> None:
        """Generate a report.

        Parameters
        ----------
        simulation : Simulation
            The Simulation to generate a report for
        state : State
            The current state of the simulation
        """
        pass

__init__ 

__init__(
    report_interval: int,
    openmm_selection: Sequence[str] = ("CA",),
) -> None

Initialize the reporter.

Parameters:

Name Type Description Default
report_interval int

The interval at which to write frames.

 required
openmm_selection Sequence[str]

The OpenMM selection strings for the atoms to use when reporting positions (default is ('CA',)).

 ('CA',)
Source code in deepdrivewe/simulation/openmm.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
def __init__(
    self,
    report_interval: int,
    openmm_selection: Sequence[str] = ('CA',),
) -> None:
    """Initialize the reporter.

    Parameters
    ----------
    report_interval : int
        The interval at which to write frames.
    openmm_selection : Sequence[str]
        The OpenMM selection strings for the atoms to use
        when reporting positions (default is ('CA',)).
    """
    self.report_interval = report_interval
    self.openmm_selection = openmm_selection

describeNextReport 

describeNextReport(
    simulation: Simulation,
) -> tuple[int, bool, bool, bool, bool, bool | None]

Get information about the next report this object will generate.

Parameters:

Name Type Description Default
simulation Simulation

The Simulation to generate a report for

 required

Returns:

Type Description
tuple

A six element tuple. The first element is the number of steps until the next report. The next four elements specify whether that report will require positions, velocities, forces, and energies respectively. The final element specifies whether positions should be wrapped to lie in a single periodic box.
Source code in deepdrivewe/simulation/openmm.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
def describeNextReport(  # noqa: N802
    self,
    simulation: app.Simulation,
) -> tuple[int, bool, bool, bool, bool, bool | None]:
    """Get information about the next report this object will generate.

    Parameters
    ----------
    simulation : Simulation
        The Simulation to generate a report for

    Returns
    -------
    tuple
        A six element tuple. The first element is the number of steps
        until the next report. The next four elements specify whether
        that report will require positions, velocities, forces, and
        energies respectively.  The final element specifies whether
        positions should be wrapped to lie in a single periodic box.
    """
    step_progress = simulation.currentStep % self.report_interval
    steps = self.report_interval - step_progress
    return (steps, True, False, False, False, None)

get_positions 

get_positions(
    simulation: Simulation, state: State
) -> np.ndarray

Get the atomic positions from the simulation.

Parameters:

Name Type Description Default
simulation Simulation

The Simulation to generate a report for.

 required
state State

The current state of the simulation.

 required

Returns:

Type Description
ndarray

The atomic positions from the simulation.
Source code in deepdrivewe/simulation/openmm.py 
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
def get_positions(
    self,
    simulation: app.Simulation,
    state: openmm.State,
) -> np.ndarray:
    """Get the atomic positions from the simulation.

    Parameters
    ----------
    simulation : Simulation
        The Simulation to generate a report for.
    state : State
        The current state of the simulation.

    Returns
    -------
    np.ndarray
        The atomic positions from the simulation.
    """
    # Get the atom indices for the selection
    atom_indices = [
        a.index
        for a in simulation.topology.atoms()
        if a.name in self.openmm_selection
    ]

    # Get the atomic coordinates of the selection
    positions = state.getPositions(asNumpy=True)
    positions = positions[atom_indices].astype(np.float32)

    # Convert positions from nanometers to angstroms
    positions *= 10.0

    return positions

report abstractmethod 

report(simulation: Simulation, state: State) -> None

Generate a report.

Parameters:

Name Type Description Default
simulation Simulation

The Simulation to generate a report for

 required
state State

The current state of the simulation

 required
Source code in deepdrivewe/simulation/openmm.py 
119
120
121
122
123
124
125
126
127
128
129
130
@abstractmethod
def report(self, simulation: app.Simulation, state: openmm.State) -> None:
    """Generate a report.

    Parameters
    ----------
    simulation : Simulation
        The Simulation to generate a report for
    state : State
        The current state of the simulation
    """
    pass

Collector

Bases: ABC

Collector interface for OpenMM simulations.

Source code in deepdrivewe/simulation/openmm.py 
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
class Collector(ABC):
    """Collector interface for OpenMM simulations."""

    def __init__(self, topic: str) -> None:
        """Initialize the collector."""
        self._topic = topic

    @property
    def topic(self) -> str:
        """The topic of the collector (a string identifier)."""
        return self._topic

    @abstractmethod
    def get(self) -> np.ndarray:
        """Get the collected data from the simulation.

        Returns
        -------
        np.ndarray
            The collected data from the simulation.
        """
        ...

    @abstractmethod
    def collect(self, positions: np.ndarray) -> Any:
        """Collect data from the simulation.

        Parameters
        ----------
        positions : np.ndarray
            The atomic positions from the simulation.
        """
        ...

topic property 

topic: str

The topic of the collector (a string identifier).

__init__ 

__init__(topic: str) -> None

Initialize the collector.

Source code in deepdrivewe/simulation/openmm.py 
136
137
138
def __init__(self, topic: str) -> None:
    """Initialize the collector."""
    self._topic = topic

get abstractmethod 

get() -> np.ndarray

Get the collected data from the simulation.

Returns:

Type Description
ndarray

The collected data from the simulation.
Source code in deepdrivewe/simulation/openmm.py 
145
146
147
148
149
150
151
152
153
154
@abstractmethod
def get(self) -> np.ndarray:
    """Get the collected data from the simulation.

    Returns
    -------
    np.ndarray
        The collected data from the simulation.
    """
    ...

collect abstractmethod 

collect(positions: ndarray) -> Any

Collect data from the simulation.

Parameters:

Name Type Description Default
positions ndarray

The atomic positions from the simulation.

 required
Source code in deepdrivewe/simulation/openmm.py 
156
157
158
159
160
161
162
163
164
165
@abstractmethod
def collect(self, positions: np.ndarray) -> Any:
    """Collect data from the simulation.

    Parameters
    ----------
    positions : np.ndarray
        The atomic positions from the simulation.
    """
    ...

CoordinatesCollector

Bases: Collector

Coordinates collector for OpenMM simulations.

Source code in deepdrivewe/simulation/openmm.py 
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
class CoordinatesCollector(Collector):
    """Coordinates collector for OpenMM simulations."""

    def __init__(
        self,
        reference_file: Path | None = None,
        mda_selection: str = 'protein and name CA',
        topic: str = 'coordinates',
    ) -> None:
        """Initialize the coordinates collector.

        Parameters
        ----------
        reference_file : Path | None
            The reference PDB file for the analysis (if provided,
            the coordinates will be aligned to the reference).
            Default is None.
        mda_selection : str
            The MDAnalysis selection string for the atoms to use
            for alignment (default is 'protein and name CA').
        topic : str
            The topic of the collector, default is 'coordinates'.
        """
        super().__init__(topic)
        self._coordinates: list[np.ndarray] = []
        self._ref = None

        # If provided, load the reference structure and save the positions
        if reference_file is not None:
            mda_u = MDAnalysis.Universe(reference_file)
            self._ref = mda_u.select_atoms(mda_selection).positions.copy()

    def get(self) -> np.ndarray:
        """Get the coordinates from the simulation.

        Returns
        -------
        np.ndarray
            The atomic positions from each frame of the simulation
            (n_frames, n_atoms, 3). Where n_atoms is the number of atoms
            in the openmm_selection of the corresponding OpenMMReporter.
        """
        return np.array(self._coordinates)

    def _align(self, positions: np.ndarray) -> np.ndarray:
        """Align the atomic positions to the reference.

        Parameters
        ----------
        positions : np.ndarray
            The atomic positions from the simulation.

        Returns
        -------
        np.ndarray
            The aligned atomic positions.
        """
        # Calculate rotation and translation using align.rotation_matrix()
        rotation_matrix, _ = align.rotation_matrix(positions, self._ref)

        # Apply the rotation to the raw positions
        aligned_positions = np.dot(positions, rotation_matrix.T)

        return aligned_positions

    def collect(self, positions: np.ndarray) -> np.ndarray:
        """Generate a report.

        Parameters
        ----------
        positions : np.ndarray
            The atomic positions from the simulation.
        """
        # Align the coordinates to the reference if provided
        pos = positions.copy() if self._ref is None else self._align(positions)

        # Collect the position coordinates
        self._coordinates.append(pos)

        return pos

__init__ 

__init__(
    reference_file: Path | None = None,
    mda_selection: str = "protein and name CA",
    topic: str = "coordinates",
) -> None

Initialize the coordinates collector.

Parameters:

Name Type Description Default
reference_file Path | None

The reference PDB file for the analysis (if provided, the coordinates will be aligned to the reference). Default is None.

 None
mda_selection str

The MDAnalysis selection string for the atoms to use for alignment (default is 'protein and name CA').

 'protein and name CA'
topic str

The topic of the collector, default is 'coordinates'.

 'coordinates'
Source code in deepdrivewe/simulation/openmm.py 
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
def __init__(
    self,
    reference_file: Path | None = None,
    mda_selection: str = 'protein and name CA',
    topic: str = 'coordinates',
) -> None:
    """Initialize the coordinates collector.

    Parameters
    ----------
    reference_file : Path | None
        The reference PDB file for the analysis (if provided,
        the coordinates will be aligned to the reference).
        Default is None.
    mda_selection : str
        The MDAnalysis selection string for the atoms to use
        for alignment (default is 'protein and name CA').
    topic : str
        The topic of the collector, default is 'coordinates'.
    """
    super().__init__(topic)
    self._coordinates: list[np.ndarray] = []
    self._ref = None

    # If provided, load the reference structure and save the positions
    if reference_file is not None:
        mda_u = MDAnalysis.Universe(reference_file)
        self._ref = mda_u.select_atoms(mda_selection).positions.copy()

get 

get() -> np.ndarray

Get the coordinates from the simulation.

Returns:

Type Description
ndarray

The atomic positions from each frame of the simulation (n_frames, n_atoms, 3). Where n_atoms is the number of atoms in the openmm_selection of the corresponding OpenMMReporter.
Source code in deepdrivewe/simulation/openmm.py 
200
201
202
203
204
205
206
207
208
209
210
def get(self) -> np.ndarray:
    """Get the coordinates from the simulation.

    Returns
    -------
    np.ndarray
        The atomic positions from each frame of the simulation
        (n_frames, n_atoms, 3). Where n_atoms is the number of atoms
        in the openmm_selection of the corresponding OpenMMReporter.
    """
    return np.array(self._coordinates)

collect 

collect(positions: ndarray) -> np.ndarray

Generate a report.

Parameters:

Name Type Description Default
positions ndarray

The atomic positions from the simulation.

 required
Source code in deepdrivewe/simulation/openmm.py 
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
def collect(self, positions: np.ndarray) -> np.ndarray:
    """Generate a report.

    Parameters
    ----------
    positions : np.ndarray
        The atomic positions from the simulation.
    """
    # Align the coordinates to the reference if provided
    pos = positions.copy() if self._ref is None else self._align(positions)

    # Collect the position coordinates
    self._coordinates.append(pos)

    return pos

RMSDCollector

Bases: Collector

RMSD collector for OpenMM simulations.

Source code in deepdrivewe/simulation/openmm.py 
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
class RMSDCollector(Collector):
    """RMSD collector for OpenMM simulations."""

    def __init__(
        self,
        reference_file: Path,
        mda_selection: str = 'protein and name CA',
        topic: str = 'rmsds',
    ) -> None:
        """Initialize the RMSD collector.

        Parameters
        ----------
        reference_file : Path
            The reference PDB file for the analysis.
        mda_selection : str
            The MDAnalysis selection string for the atoms to use
            (default is 'protein and name CA').
        topic : str
            The topic of the collector, default is 'rmsd'.
        """
        super().__init__(topic)
        self._rmsd: list[float] = []

        # Load the reference structure and save the positions
        mda_u = MDAnalysis.Universe(reference_file)
        self._ref = mda_u.select_atoms(mda_selection).positions.copy()

    def get(self) -> np.ndarray:
        """Get the RMSDs from the simulation.

        Returns
        -------
        np.ndarray
            The RMSDs from the simulation shaped as (n_frames, 1).
        """
        return np.array(self._rmsd).reshape(-1, 1)

    def collect(self, positions: np.ndarray) -> float:
        """Generate a report.

        Parameters
        ----------
        positions : np.ndarray
            The atomic positions from the simulation.
        """
        # Compute the RMSD
        rmsd = rms.rmsd(positions, self._ref, superposition=True)
        self._rmsd.append(rmsd)
        return rmsd

__init__ 

__init__(
    reference_file: Path,
    mda_selection: str = "protein and name CA",
    topic: str = "rmsds",
) -> None

Initialize the RMSD collector.

Parameters:

Name Type Description Default
reference_file Path

The reference PDB file for the analysis.

 required
mda_selection str

The MDAnalysis selection string for the atoms to use (default is 'protein and name CA').

 'protein and name CA'
topic str

The topic of the collector, default is 'rmsd'.

 'rmsds'
Source code in deepdrivewe/simulation/openmm.py 
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
def __init__(
    self,
    reference_file: Path,
    mda_selection: str = 'protein and name CA',
    topic: str = 'rmsds',
) -> None:
    """Initialize the RMSD collector.

    Parameters
    ----------
    reference_file : Path
        The reference PDB file for the analysis.
    mda_selection : str
        The MDAnalysis selection string for the atoms to use
        (default is 'protein and name CA').
    topic : str
        The topic of the collector, default is 'rmsd'.
    """
    super().__init__(topic)
    self._rmsd: list[float] = []

    # Load the reference structure and save the positions
    mda_u = MDAnalysis.Universe(reference_file)
    self._ref = mda_u.select_atoms(mda_selection).positions.copy()

get 

get() -> np.ndarray

Get the RMSDs from the simulation.

Returns:

Type Description
ndarray

The RMSDs from the simulation shaped as (n_frames, 1).
Source code in deepdrivewe/simulation/openmm.py 
278
279
280
281
282
283
284
285
286
def get(self) -> np.ndarray:
    """Get the RMSDs from the simulation.

    Returns
    -------
    np.ndarray
        The RMSDs from the simulation shaped as (n_frames, 1).
    """
    return np.array(self._rmsd).reshape(-1, 1)

collect 

collect(positions: ndarray) -> float

Generate a report.

Parameters:

Name Type Description Default
positions ndarray

The atomic positions from the simulation.

 required
Source code in deepdrivewe/simulation/openmm.py 
288
289
290
291
292
293
294
295
296
297
298
299
def collect(self, positions: np.ndarray) -> float:
    """Generate a report.

    Parameters
    ----------
    positions : np.ndarray
        The atomic positions from the simulation.
    """
    # Compute the RMSD
    rmsd = rms.rmsd(positions, self._ref, superposition=True)
    self._rmsd.append(rmsd)
    return rmsd

ContactMapCollector

Bases: Collector

Contact map collector for OpenMM simulations.

Source code in deepdrivewe/simulation/openmm.py 
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
class ContactMapCollector(Collector):
    """Contact map collector for OpenMM simulations."""

    def __init__(
        self,
        cutoff_angstrom: float = 8.0,
        topic: str = 'contact_maps',
    ) -> None:
        """Initialize the contact map collector."""
        super().__init__(topic)
        self.cutoff_angstrom = cutoff_angstrom
        self._contact_maps: list[np.ndarray] = []

    def get(self) -> np.ndarray:
        """Get the contact maps from the simulation.

        Returns
        -------
        np.ndarray
            The contact maps from the simulation as a ragged array
            shaped as (n_frames, *).
        """
        # Collect the contact maps in a ragged numpy array
        contact_maps = np.array(self._contact_maps, dtype=object)

        return contact_maps

    def collect(self, positions: np.ndarray) -> np.ndarray:
        """Generate a report.

        Parameters
        ----------
        positions : np.ndarray
            The atomic positions from the simulation.
        """
        # Compute the contact map
        contact_map = distances.contact_matrix(
            positions,
            self.cutoff_angstrom,
            returntype='sparse',
        )

        # Convert the contact map to sparse format
        coo_matrix = contact_map.tocoo()

        # Get the row and col indices and concatenate them
        row = coo_matrix.row.astype('int16')
        col = coo_matrix.col.astype('int16')
        sparse_contact_map = np.concatenate([row, col])

        # Append the row and col indices to lists
        self._contact_maps.append(sparse_contact_map)

        return sparse_contact_map

__init__ 

__init__(
    cutoff_angstrom: float = 8.0,
    topic: str = "contact_maps",
) -> None

Initialize the contact map collector.

Source code in deepdrivewe/simulation/openmm.py 
305
306
307
308
309
310
311
312
313
def __init__(
    self,
    cutoff_angstrom: float = 8.0,
    topic: str = 'contact_maps',
) -> None:
    """Initialize the contact map collector."""
    super().__init__(topic)
    self.cutoff_angstrom = cutoff_angstrom
    self._contact_maps: list[np.ndarray] = []

get 

get() -> np.ndarray

Get the contact maps from the simulation.

Returns:

Type Description
ndarray

The contact maps from the simulation as a ragged array shaped as (n_frames, *).
Source code in deepdrivewe/simulation/openmm.py 
315
316
317
318
319
320
321
322
323
324
325
326
327
def get(self) -> np.ndarray:
    """Get the contact maps from the simulation.

    Returns
    -------
    np.ndarray
        The contact maps from the simulation as a ragged array
        shaped as (n_frames, *).
    """
    # Collect the contact maps in a ragged numpy array
    contact_maps = np.array(self._contact_maps, dtype=object)

    return contact_maps

collect 

collect(positions: ndarray) -> np.ndarray

Generate a report.

Parameters:

Name Type Description Default
positions ndarray

The atomic positions from the simulation.

 required
Source code in deepdrivewe/simulation/openmm.py 
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
def collect(self, positions: np.ndarray) -> np.ndarray:
    """Generate a report.

    Parameters
    ----------
    positions : np.ndarray
        The atomic positions from the simulation.
    """
    # Compute the contact map
    contact_map = distances.contact_matrix(
        positions,
        self.cutoff_angstrom,
        returntype='sparse',
    )

    # Convert the contact map to sparse format
    coo_matrix = contact_map.tocoo()

    # Get the row and col indices and concatenate them
    row = coo_matrix.row.astype('int16')
    col = coo_matrix.col.astype('int16')
    sparse_contact_map = np.concatenate([row, col])

    # Append the row and col indices to lists
    self._contact_maps.append(sparse_contact_map)

    return sparse_contact_map

CollectionReporter

Bases: OpenMMReporter

Reporter to collect multiple data products from an OpenMM simulation.

Source code in deepdrivewe/simulation/openmm.py 
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
class CollectionReporter(OpenMMReporter):
    """Reporter to collect multiple data products from an OpenMM simulation."""

    def __init__(
        self,
        report_interval: int,
        collectors: list[Collector],
        openmm_selection: Sequence[str] = ('CA',),
    ) -> None:
        """Initialize the reporter.

        Parameters
        ----------
        report_interval : int
            The interval at which to write frames.
        collectors : list[Collector]
            The collectors to inject into the simulation.
        openmm_selection : Sequence[str]
            The OpenMM selection strings for the atoms to use
            when reporting positions (default is ('CA',)).

        Raises
        ------
        ValueError
            If the collectors have duplicate topics.
        """
        super().__init__(report_interval, openmm_selection)

        # Check that the collectors have unique topics
        if len(collectors) != len({x.topic for x in collectors}):
            raise ValueError('Collectors must have unique topics.')

        self.collectors = collectors

    def get_collected_data(self) -> dict[str, np.ndarray]:
        """Get the collected data from the simulation.

        Returns
        -------
        dict[str, np.ndarray]
            The collected data from the simulation.
        """
        return {x.topic: x.get() for x in self.collectors}

    def report(self, simulation: app.Simulation, state: openmm.State) -> None:
        """Generate a report.

        Parameters
        ----------
        simulation : Simulation
            The Simulation to generate a report for.
        state : State
            The current state of the simulation.
        """
        # Get the positions
        positions = self.get_positions(simulation, state)

        # Collect data from the simulation
        for collector in self.collectors:
            collector.collect(positions)

__init__ 

__init__(
    report_interval: int,
    collectors: list[Collector],
    openmm_selection: Sequence[str] = ("CA",),
) -> None

Initialize the reporter.

Parameters:

Name Type Description Default
report_interval int

The interval at which to write frames.

 required
collectors list[Collector]

The collectors to inject into the simulation.

 required
openmm_selection Sequence[str]

The OpenMM selection strings for the atoms to use when reporting positions (default is ('CA',)).

 ('CA',)

Raises:

Type Description
ValueError

If the collectors have duplicate topics.
Source code in deepdrivewe/simulation/openmm.py 
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
def __init__(
    self,
    report_interval: int,
    collectors: list[Collector],
    openmm_selection: Sequence[str] = ('CA',),
) -> None:
    """Initialize the reporter.

    Parameters
    ----------
    report_interval : int
        The interval at which to write frames.
    collectors : list[Collector]
        The collectors to inject into the simulation.
    openmm_selection : Sequence[str]
        The OpenMM selection strings for the atoms to use
        when reporting positions (default is ('CA',)).

    Raises
    ------
    ValueError
        If the collectors have duplicate topics.
    """
    super().__init__(report_interval, openmm_selection)

    # Check that the collectors have unique topics
    if len(collectors) != len({x.topic for x in collectors}):
        raise ValueError('Collectors must have unique topics.')

    self.collectors = collectors

get_collected_data 

get_collected_data() -> dict[str, np.ndarray]

Get the collected data from the simulation.

Returns:

Type Description
dict[str, ndarray]

The collected data from the simulation.
Source code in deepdrivewe/simulation/openmm.py 
392
393
394
395
396
397
398
399
400
def get_collected_data(self) -> dict[str, np.ndarray]:
    """Get the collected data from the simulation.

    Returns
    -------
    dict[str, np.ndarray]
        The collected data from the simulation.
    """
    return {x.topic: x.get() for x in self.collectors}

report 

report(simulation: Simulation, state: State) -> None

Generate a report.

Parameters:

Name Type Description Default
simulation Simulation

The Simulation to generate a report for.

 required
state State

The current state of the simulation.

 required
Source code in deepdrivewe/simulation/openmm.py 
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
def report(self, simulation: app.Simulation, state: openmm.State) -> None:
    """Generate a report.

    Parameters
    ----------
    simulation : Simulation
        The Simulation to generate a report for.
    state : State
        The current state of the simulation.
    """
    # Get the positions
    positions = self.get_positions(simulation, state)

    # Collect data from the simulation
    for collector in self.collectors:
        collector.collect(positions)

OpenMMConfig

Bases: BaseModel

Configuration for an OpenMM simulation.

Source code in deepdrivewe/simulation/openmm.py 
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
class OpenMMConfig(BaseModel):
    """Configuration for an OpenMM simulation."""

    simulation_length_ns: float = Field(
        default=0.01,  # 0.01 ns = 10 ps
        description='The length of the simulation (in nanoseconds).',
    )
    report_interval_ps: float = Field(
        default=2.0,
        description='The report interval for the simulation (in picoseconds).',
    )
    dt_ps: float = Field(
        default=0.002,
        description='The timestep for the simulation.',
    )
    temperature_kelvin: float = Field(
        default=300.0,
        description='The temperature for the simulation.',
    )
    heat_bath_friction_coef: float = Field(
        default=1.0,
        description='The heat bath friction coefficient for the simulation.',
    )
    solvent_type: str = Field(
        default='implicit',
        description='The solvent type for the simulation.',
    )
    explicit_barostat: str | None = Field(
        default=None,
        description='The barostat used for an explicit solvent simulation. '
        'Options are: None (NVT), MonteCarloBarostat, ,'
        'MonteCarloAnisotropicBarostat.',
    )
    run_minimization: bool = Field(
        default=True,
        description='Whether to run energy minimization.',
    )
    set_positions: bool = Field(
        default=True,
        description='Whether to set positions.',
    )
    randomize_velocities: bool = Field(
        default=True,
        description='Whether to randomize the basis state initial velocities.',
    )
    hardware_platform: str = Field(
        default='CUDA',
        description='The hardware platform to use for the simulation.'
        ' Options are: CUDA, OpenCL, CPU.',
    )

    @model_validator(mode='after')
    def validate_explicit_barostat(self) -> Self:
        """Check for valid explicit_barostat options."""
        valid_barostats = (
            None,
            'MonteCarloBarostat',
            'MonteCarloAnisotropicBarostat',
        )
        if (
            self.solvent_type == 'explicit'
            and self.explicit_barostat not in valid_barostats
        ):
            raise ValueError(
                f'Invalid explicit_barostat option: {self.explicit_barostat}',
                f'For explicit solvent, valid options are: {valid_barostats}',
            )
        return self

    @model_validator(mode='after')
    def validate_hardware_platform(self) -> Self:
        """Check for valid hardware_platform options."""
        valid_platforms = ('CUDA', 'OpenCL', 'CPU')
        if self.hardware_platform not in valid_platforms:
            raise ValueError(
                f'Invalid hardware_platform option: {self.hardware_platform}',
                f'Valid options are: {valid_platforms}',
            )
        return self

    @property
    def num_steps(self) -> int:
        """The number of steps to run the simulation."""
        dt_ps = self.dt_ps * u.picoseconds
        simulation_length_ns = self.simulation_length_ns * u.nanoseconds
        return int(simulation_length_ns / dt_ps)

    @property
    def report_steps(self) -> int:
        """The number of steps between log reports."""
        dt_ps = self.dt_ps * u.picoseconds
        report_interval_ps = self.report_interval_ps * u.picoseconds
        return int(report_interval_ps / dt_ps)

    def load_explicit_system_from_top(
        self,
        top_file: str | Path,
        pdb_file: str | Path,
    ) -> tuple[openmm.System, app.Topology]:
        """Load an explicit solvent system from a topology file.

        Parameters
        ----------
        top_file : str | Path
            The topology file to load the system from.
        pdb_file : str | Path
            The PDB file to load the system topology.

        Returns
        -------
        tuple[openmm.System, app.Topology]
            The OpenMM system and topology.
        """
        # Load the topology file
        top = pmd.load_file(str(top_file), str(pdb_file))

        # Configure system
        system = top.createSystem(
            nonbondedMethod=app.PME,
            nonbondedCutoff=1.0 * u.nanometer,
            constraints=app.HBonds,
        )

        return system, top.topology

    def load_implicit_system_from_pdb(
        self,
        pdb_file: str | Path,
    ) -> tuple[openmm.System, app.Topology]:
        """Load an implicit solvent system from a PDB file.

        Parameters
        ----------
        pdb_file : str | Path
            The PDB file to load the system from.

        Returns
        -------
        tuple[openmm.System, app.Topology]
            The OpenMM system and topology.
        """
        # Load the PDB file
        pdb = app.PDBFile(str(pdb_file))

        # Get the topology
        topology = pdb.topology

        # Set the forcefield
        forcefield = app.ForceField('amber14-all.xml', 'implicit/gbn2.xml')

        # Configure the system
        system = forcefield.createSystem(
            topology,
            nonbondedMethod=app.CutoffNonPeriodic,
            nonbondedCutoff=1.0 * u.nanometer,
            constraints=app.HBonds,
        )

        return system, topology

    def load_implicit_system_from_top(
        self,
        top_file: str | Path,
    ) -> tuple[openmm.System, app.Topology]:
        """Load an implicit solvent system from a topology file.

        Parameters
        ----------
        top_file : str | Path
            The topology file to load the system from.

        Returns
        -------
        tuple[openmm.System, app.Topology]
            The OpenMM system and topology.
        """
        # Load the topology file
        top = app.AmberPrmtopFile(str(top_file))

        # Configure the system
        system = top.createSystem(
            nonbondedMethod=app.CutoffNonPeriodic,
            nonbondedCutoff=1.0 * u.nanometer,
            constraints=app.HBonds,
            implicitSolvent=app.OBC1,
        )

        return system, top.topology

    def configure_hardware(self) -> tuple[openmm.Platform, dict[str, str]]:
        """Configure the hardware for the simulation.

        Returns
        -------
        tuple[openmm.Platform, dict[str, str]]
            The OpenMM platform and the platform properties.
        """
        if self.hardware_platform == 'CUDA':
            # Use the CUDA platform
            platform = openmm.Platform.getPlatformByName('CUDA')
            platform_properties = {
                'DeviceIndex': '0',
                'CudaPrecision': 'mixed',
            }
        elif self.hardware_platform == 'OpenCL':
            # Use the OpenCL platform
            platform = openmm.Platform.getPlatformByName('OpenCL')
            platform_properties = {'DeviceIndex': '0'}
        else:
            # Use the CPU platform
            platform = openmm.Platform.getPlatformByName('CPU')
            platform_properties = {}

        return platform, platform_properties

    def configure_integrator(self) -> openmm.LangevinIntegrator:
        """Configure the integrator for the simulation.

        Returns
        -------
        openmm.LangevinIntegrator
            The configured integrator.
        """
        # Configure the integrator
        integrator = openmm.LangevinIntegrator(
            self.temperature_kelvin * u.kelvin,
            self.heat_bath_friction_coef / u.picosecond,
            self.dt_ps * u.picosecond,
        )

        # Set the constraint tolerance
        integrator.setConstraintTolerance(0.00001)

        return integrator

    def configure_barostat(
        self,
    ) -> (
        openmm.MonteCarloBarostat | openmm.MonteCarloAnisotropicBarostat | None
    ):
        """Configure the barostat for the simulation.

        Returns
        -------
        openmm.MonteCarloBarostat | openmm.MonteCarloAnisotropicBarostat | None
            The configured barostat or None if no barostat is used.
        """
        if self.explicit_barostat == 'MonteCarloBarostat':
            return openmm.MonteCarloBarostat(
                1 * u.bar,
                self.temperature_kelvin * u.kelvin,
            )

        elif self.explicit_barostat == 'MonteCarloAnisotropicBarostat':
            return openmm.MonteCarloAnisotropicBarostat(
                (1, 1, 1) * u.bar,
                self.temperature_kelvin * u.kelvin,
                False,
                False,
                True,
            )

        return None

    # NOTE: Add a retry decorator to the method since sometimes the
    # PDB file is not fully written before it is read.
    @retry_on_exception(wait_time=30)
    def configure_simulation(
        self,
        pdb_file: str | Path,
        top_file: str | Path | None,
        checkpoint_file: str | Path | None = None,
    ) -> app.Simulation:
        """Configure an OpenMM simulation.

        Parameters
        ----------
        pdb_file : str | Path
            The PDB file to initialize the positions (used to load the system
            topology for implicit solvent).
        top_file : str | Path | None
            The optional topology file to initialize the systems topology
            (required for explicit solvent).
        checkpoint_file : str | Path | None
            The checkpoint file to initialize the simulation.

        Returns
        -------
        app.Simulation
            Configured OpenMM Simulation object.

        Raises
        ------
        ValueError
            If explicit solvent is selected and no topology file is provided.
        """
        # Select implicit or explicit solvent configuration and load the system
        if self.solvent_type == 'explicit':
            if top_file is None:
                raise ValueError(
                    'Topology file must be provided for explicit solvent.',
                )
            system, topology = self.load_explicit_system_from_top(
                pdb_file,
                top_file,
            )
        elif top_file is not None:
            system, topology = self.load_implicit_system_from_top(top_file)
        else:
            system, topology = self.load_implicit_system_from_pdb(pdb_file)

        # Configure the integrator
        integrator = self.configure_integrator()

        # Configure the barostat
        barostat = self.configure_barostat()
        if barostat is not None:
            system.addForce(barostat)

        # Configure the hardware
        platform, platform_properties = self.configure_hardware()

        # Create the simulation
        sim = app.Simulation(
            topology,
            system,
            integrator,
            platform,
            platform_properties,
        )

        # Load the checkpoint file if provided (skips setting positions
        # from PDB, minimization, and randomizing velocities)
        if checkpoint_file is not None:
            # Load the checkpoint file
            sim.loadCheckpoint(str(checkpoint_file))

            # Create a new Simulation with the existing system context,
            # but with the new integrator (which applies the new RNG seed)
            new_simulation = app.Simulation(
                sim.topology,
                sim.system,
                self.configure_integrator(),
                platform,
                platform_properties,
                # sim.context.getPlatform(),
                # *self.configure_hardware(),
            )

            # Set the state from the existing context to continue the sim
            state = sim.context.getState(
                getPositions=True,
                getVelocities=True,
                getEnergy=True,
                getForces=True,
            )
            new_simulation.context.setState(state)

            return new_simulation

        # Set the positions
        if self.set_positions:
            pdb = app.PDBFile(str(pdb_file))
            sim.context.setPositions(pdb.getPositions())

        # Minimize energy and equilibrate
        if self.run_minimization:
            sim.minimizeEnergy()

        # Set velocities to temperature
        if self.randomize_velocities:
            sim.context.setVelocitiesToTemperature(
                self.temperature_kelvin * u.kelvin,
                random.randint(1, 10000),
            )

        return sim

num_steps property 

num_steps: int

The number of steps to run the simulation.

report_steps property 

report_steps: int

The number of steps between log reports.

validate_explicit_barostat 

validate_explicit_barostat() -> Self

Check for valid explicit_barostat options.

Source code in deepdrivewe/simulation/openmm.py 
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
@model_validator(mode='after')
def validate_explicit_barostat(self) -> Self:
    """Check for valid explicit_barostat options."""
    valid_barostats = (
        None,
        'MonteCarloBarostat',
        'MonteCarloAnisotropicBarostat',
    )
    if (
        self.solvent_type == 'explicit'
        and self.explicit_barostat not in valid_barostats
    ):
        raise ValueError(
            f'Invalid explicit_barostat option: {self.explicit_barostat}',
            f'For explicit solvent, valid options are: {valid_barostats}',
        )
    return self

validate_hardware_platform 

validate_hardware_platform() -> Self

Check for valid hardware_platform options.

Source code in deepdrivewe/simulation/openmm.py 
489
490
491
492
493
494
495
496
497
498
@model_validator(mode='after')
def validate_hardware_platform(self) -> Self:
    """Check for valid hardware_platform options."""
    valid_platforms = ('CUDA', 'OpenCL', 'CPU')
    if self.hardware_platform not in valid_platforms:
        raise ValueError(
            f'Invalid hardware_platform option: {self.hardware_platform}',
            f'Valid options are: {valid_platforms}',
        )
    return self

load_explicit_system_from_top 

load_explicit_system_from_top(
    top_file: str | Path, pdb_file: str | Path
) -> tuple[openmm.System, app.Topology]

Load an explicit solvent system from a topology file.

Parameters:

Name Type Description Default
top_file str | Path

The topology file to load the system from.

 required
pdb_file str | Path

The PDB file to load the system topology.

 required

Returns:

Type Description
tuple[System, Topology]

The OpenMM system and topology.
Source code in deepdrivewe/simulation/openmm.py 
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
def load_explicit_system_from_top(
    self,
    top_file: str | Path,
    pdb_file: str | Path,
) -> tuple[openmm.System, app.Topology]:
    """Load an explicit solvent system from a topology file.

    Parameters
    ----------
    top_file : str | Path
        The topology file to load the system from.
    pdb_file : str | Path
        The PDB file to load the system topology.

    Returns
    -------
    tuple[openmm.System, app.Topology]
        The OpenMM system and topology.
    """
    # Load the topology file
    top = pmd.load_file(str(top_file), str(pdb_file))

    # Configure system
    system = top.createSystem(
        nonbondedMethod=app.PME,
        nonbondedCutoff=1.0 * u.nanometer,
        constraints=app.HBonds,
    )

    return system, top.topology

load_implicit_system_from_pdb 

load_implicit_system_from_pdb(
    pdb_file: str | Path,
) -> tuple[openmm.System, app.Topology]

Load an implicit solvent system from a PDB file.

Parameters:

Name Type Description Default
pdb_file str | Path

The PDB file to load the system from.

 required

Returns:

Type Description
tuple[System, Topology]

The OpenMM system and topology.
Source code in deepdrivewe/simulation/openmm.py 
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
def load_implicit_system_from_pdb(
    self,
    pdb_file: str | Path,
) -> tuple[openmm.System, app.Topology]:
    """Load an implicit solvent system from a PDB file.

    Parameters
    ----------
    pdb_file : str | Path
        The PDB file to load the system from.

    Returns
    -------
    tuple[openmm.System, app.Topology]
        The OpenMM system and topology.
    """
    # Load the PDB file
    pdb = app.PDBFile(str(pdb_file))

    # Get the topology
    topology = pdb.topology

    # Set the forcefield
    forcefield = app.ForceField('amber14-all.xml', 'implicit/gbn2.xml')

    # Configure the system
    system = forcefield.createSystem(
        topology,
        nonbondedMethod=app.CutoffNonPeriodic,
        nonbondedCutoff=1.0 * u.nanometer,
        constraints=app.HBonds,
    )

    return system, topology

load_implicit_system_from_top 

load_implicit_system_from_top(
    top_file: str | Path,
) -> tuple[openmm.System, app.Topology]

Load an implicit solvent system from a topology file.

Parameters:

Name Type Description Default
top_file str | Path

The topology file to load the system from.

 required

Returns:

Type Description
tuple[System, Topology]

The OpenMM system and topology.
Source code in deepdrivewe/simulation/openmm.py 
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
def load_implicit_system_from_top(
    self,
    top_file: str | Path,
) -> tuple[openmm.System, app.Topology]:
    """Load an implicit solvent system from a topology file.

    Parameters
    ----------
    top_file : str | Path
        The topology file to load the system from.

    Returns
    -------
    tuple[openmm.System, app.Topology]
        The OpenMM system and topology.
    """
    # Load the topology file
    top = app.AmberPrmtopFile(str(top_file))

    # Configure the system
    system = top.createSystem(
        nonbondedMethod=app.CutoffNonPeriodic,
        nonbondedCutoff=1.0 * u.nanometer,
        constraints=app.HBonds,
        implicitSolvent=app.OBC1,
    )

    return system, top.topology

configure_hardware 

configure_hardware() -> (
    tuple[openmm.Platform, dict[str, str]]
)

Configure the hardware for the simulation.

Returns:

Type Description
tuple[Platform, dict[str, str]]

The OpenMM platform and the platform properties.
Source code in deepdrivewe/simulation/openmm.py 
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
def configure_hardware(self) -> tuple[openmm.Platform, dict[str, str]]:
    """Configure the hardware for the simulation.

    Returns
    -------
    tuple[openmm.Platform, dict[str, str]]
        The OpenMM platform and the platform properties.
    """
    if self.hardware_platform == 'CUDA':
        # Use the CUDA platform
        platform = openmm.Platform.getPlatformByName('CUDA')
        platform_properties = {
            'DeviceIndex': '0',
            'CudaPrecision': 'mixed',
        }
    elif self.hardware_platform == 'OpenCL':
        # Use the OpenCL platform
        platform = openmm.Platform.getPlatformByName('OpenCL')
        platform_properties = {'DeviceIndex': '0'}
    else:
        # Use the CPU platform
        platform = openmm.Platform.getPlatformByName('CPU')
        platform_properties = {}

    return platform, platform_properties

configure_integrator 

configure_integrator() -> openmm.LangevinIntegrator

Configure the integrator for the simulation.

Returns:

Type Description
LangevinIntegrator

The configured integrator.
Source code in deepdrivewe/simulation/openmm.py 
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
def configure_integrator(self) -> openmm.LangevinIntegrator:
    """Configure the integrator for the simulation.

    Returns
    -------
    openmm.LangevinIntegrator
        The configured integrator.
    """
    # Configure the integrator
    integrator = openmm.LangevinIntegrator(
        self.temperature_kelvin * u.kelvin,
        self.heat_bath_friction_coef / u.picosecond,
        self.dt_ps * u.picosecond,
    )

    # Set the constraint tolerance
    integrator.setConstraintTolerance(0.00001)

    return integrator

configure_barostat 

configure_barostat() -> (
    openmm.MonteCarloBarostat
    | openmm.MonteCarloAnisotropicBarostat
    | None
)

Configure the barostat for the simulation.

Returns:

Type Description
MonteCarloBarostat | MonteCarloAnisotropicBarostat | None

The configured barostat or None if no barostat is used.
Source code in deepdrivewe/simulation/openmm.py 
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
def configure_barostat(
    self,
) -> (
    openmm.MonteCarloBarostat | openmm.MonteCarloAnisotropicBarostat | None
):
    """Configure the barostat for the simulation.

    Returns
    -------
    openmm.MonteCarloBarostat | openmm.MonteCarloAnisotropicBarostat | None
        The configured barostat or None if no barostat is used.
    """
    if self.explicit_barostat == 'MonteCarloBarostat':
        return openmm.MonteCarloBarostat(
            1 * u.bar,
            self.temperature_kelvin * u.kelvin,
        )

    elif self.explicit_barostat == 'MonteCarloAnisotropicBarostat':
        return openmm.MonteCarloAnisotropicBarostat(
            (1, 1, 1) * u.bar,
            self.temperature_kelvin * u.kelvin,
            False,
            False,
            True,
        )

    return None

configure_simulation 

configure_simulation(
    pdb_file: str | Path,
    top_file: str | Path | None,
    checkpoint_file: str | Path | None = None,
) -> app.Simulation

Configure an OpenMM simulation.

Parameters:

Name Type Description Default
pdb_file str | Path

The PDB file to initialize the positions (used to load the system topology for implicit solvent).

 required
top_file str | Path | None

The optional topology file to initialize the systems topology (required for explicit solvent).

 required
checkpoint_file str | Path | None

The checkpoint file to initialize the simulation.

 None

Returns:

Type Description
Simulation

Configured OpenMM Simulation object.

Raises:

Type Description
ValueError

If explicit solvent is selected and no topology file is provided.
Source code in deepdrivewe/simulation/openmm.py 
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
@retry_on_exception(wait_time=30)
def configure_simulation(
    self,
    pdb_file: str | Path,
    top_file: str | Path | None,
    checkpoint_file: str | Path | None = None,
) -> app.Simulation:
    """Configure an OpenMM simulation.

    Parameters
    ----------
    pdb_file : str | Path
        The PDB file to initialize the positions (used to load the system
        topology for implicit solvent).
    top_file : str | Path | None
        The optional topology file to initialize the systems topology
        (required for explicit solvent).
    checkpoint_file : str | Path | None
        The checkpoint file to initialize the simulation.

    Returns
    -------
    app.Simulation
        Configured OpenMM Simulation object.

    Raises
    ------
    ValueError
        If explicit solvent is selected and no topology file is provided.
    """
    # Select implicit or explicit solvent configuration and load the system
    if self.solvent_type == 'explicit':
        if top_file is None:
            raise ValueError(
                'Topology file must be provided for explicit solvent.',
            )
        system, topology = self.load_explicit_system_from_top(
            pdb_file,
            top_file,
        )
    elif top_file is not None:
        system, topology = self.load_implicit_system_from_top(top_file)
    else:
        system, topology = self.load_implicit_system_from_pdb(pdb_file)

    # Configure the integrator
    integrator = self.configure_integrator()

    # Configure the barostat
    barostat = self.configure_barostat()
    if barostat is not None:
        system.addForce(barostat)

    # Configure the hardware
    platform, platform_properties = self.configure_hardware()

    # Create the simulation
    sim = app.Simulation(
        topology,
        system,
        integrator,
        platform,
        platform_properties,
    )

    # Load the checkpoint file if provided (skips setting positions
    # from PDB, minimization, and randomizing velocities)
    if checkpoint_file is not None:
        # Load the checkpoint file
        sim.loadCheckpoint(str(checkpoint_file))

        # Create a new Simulation with the existing system context,
        # but with the new integrator (which applies the new RNG seed)
        new_simulation = app.Simulation(
            sim.topology,
            sim.system,
            self.configure_integrator(),
            platform,
            platform_properties,
            # sim.context.getPlatform(),
            # *self.configure_hardware(),
        )

        # Set the state from the existing context to continue the sim
        state = sim.context.getState(
            getPositions=True,
            getVelocities=True,
            getEnergy=True,
            getForces=True,
        )
        new_simulation.context.setState(state)

        return new_simulation

    # Set the positions
    if self.set_positions:
        pdb = app.PDBFile(str(pdb_file))
        sim.context.setPositions(pdb.getPositions())

    # Minimize energy and equilibrate
    if self.run_minimization:
        sim.minimizeEnergy()

    # Set velocities to temperature
    if self.randomize_velocities:
        sim.context.setVelocitiesToTemperature(
            self.temperature_kelvin * u.kelvin,
            random.randint(1, 10000),
        )

    return sim

OpenMMSimulation

Bases: BaseModel

OpenMM simulation.

Source code in deepdrivewe/simulation/openmm.py 
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
class OpenMMSimulation(BaseModel):
    """OpenMM simulation."""

    config: OpenMMConfig = Field(
        description='The configuration for the OpenMM simulation.',
    )
    checkpoint_file: Path = Field(
        description='The checkpoint file for the simulation.',
    )
    top_file: Path | None = Field(
        default=None,
        description='The topology file for the simulation.',
    )
    output_dir: Path = Field(
        description='The output directory for the simulation.',
    )
    copy_input_files: bool = Field(
        default=True,
        description='Whether to copy the input files to the output directory.'
        'top_file will be copied by default.',
    )

    @property
    def trajectory_file(self) -> Path:
        """The trajectory file for the simulation."""
        return self.output_dir / 'seg.dcd'

    @property
    def restart_file(self) -> Path:
        """The restart file for the simulation.

        NOTE: In the case of OpenMM, this PDB file is used
        to initialize the simulation, and used as a proxy
        for the restart file (which is a checkpoint file).
        The actual checkpoint file is seg.chk (found in the
        same directory as the checkpoint file) which is used
        to save the simulation state and is automatically
        loaded if it exists.
        """
        return self.output_dir / 'seg.pdb'

    @property
    def log_file(self) -> Path:
        """The log file for the simulation."""
        return self.output_dir / 'seg.log'

    def run(self, reporters: list[OpenMMReporter] | None = None) -> None:
        """Run the simulation.

        Parameters
        ----------
        reporters : list[OpenMMReporter], optional
            Custom reporters to inject into the simulation, by default None.
        """
        # Copy the restart checkpoint to the output directory
        shutil.copy(self.checkpoint_file, self.restart_file)

        # Copy the static input files to the output directory
        if self.copy_input_files and self.top_file is not None:
            self.top_file = Path(shutil.copy(self.top_file, self.output_dir))

        # Attempt to locate a checkpoint file
        chk_file = self.checkpoint_file.parent / 'seg.chk'
        checkpoint_file = chk_file if chk_file.exists() else None

        # Initialize an OpenMM simulation
        sim = self.config.configure_simulation(
            pdb_file=self.restart_file,
            top_file=self.top_file,
            checkpoint_file=checkpoint_file,
        )

        # Set up a reporter to write a simulation trajectory file
        sim.reporters.append(
            app.DCDReporter(
                file=self.trajectory_file,
                reportInterval=self.config.report_steps,
            ),
        )

        # Set up a reporter to write a simulation log file
        sim.reporters.append(
            app.StateDataReporter(
                file=str(self.log_file),
                reportInterval=self.config.report_steps,
                step=True,
                time=True,
                speed=True,
                potentialEnergy=True,
                temperature=True,
                totalEnergy=True,
            ),
        )

        # Inject the custom reporters
        if reporters is not None:
            sim.reporters.extend(reporters)

        # Run simulation
        sim.step(self.config.num_steps)

        # Save a checkpoint of the final state
        sim.saveCheckpoint(str(self.output_dir / 'seg.chk'))

trajectory_file property 

trajectory_file: Path

The trajectory file for the simulation.

restart_file property 

restart_file: Path

The restart file for the simulation.

NOTE: In the case of OpenMM, this PDB file is used to initialize the simulation, and used as a proxy for the restart file (which is a checkpoint file). The actual checkpoint file is seg.chk (found in the same directory as the checkpoint file) which is used to save the simulation state and is automatically loaded if it exists.

log_file property 

log_file: Path

The log file for the simulation.

run 

run(reporters: list[OpenMMReporter] | None = None) -> None

Run the simulation.

Parameters:

Name Type Description Default
reporters list[OpenMMReporter]

Custom reporters to inject into the simulation, by default None.

 None
Source code in deepdrivewe/simulation/openmm.py 
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
def run(self, reporters: list[OpenMMReporter] | None = None) -> None:
    """Run the simulation.

    Parameters
    ----------
    reporters : list[OpenMMReporter], optional
        Custom reporters to inject into the simulation, by default None.
    """
    # Copy the restart checkpoint to the output directory
    shutil.copy(self.checkpoint_file, self.restart_file)

    # Copy the static input files to the output directory
    if self.copy_input_files and self.top_file is not None:
        self.top_file = Path(shutil.copy(self.top_file, self.output_dir))

    # Attempt to locate a checkpoint file
    chk_file = self.checkpoint_file.parent / 'seg.chk'
    checkpoint_file = chk_file if chk_file.exists() else None

    # Initialize an OpenMM simulation
    sim = self.config.configure_simulation(
        pdb_file=self.restart_file,
        top_file=self.top_file,
        checkpoint_file=checkpoint_file,
    )

    # Set up a reporter to write a simulation trajectory file
    sim.reporters.append(
        app.DCDReporter(
            file=self.trajectory_file,
            reportInterval=self.config.report_steps,
        ),
    )

    # Set up a reporter to write a simulation log file
    sim.reporters.append(
        app.StateDataReporter(
            file=str(self.log_file),
            reportInterval=self.config.report_steps,
            step=True,
            time=True,
            speed=True,
            potentialEnergy=True,
            temperature=True,
            totalEnergy=True,
        ),
    )

    # Inject the custom reporters
    if reporters is not None:
        sim.reporters.extend(reporters)

    # Run simulation
    sim.step(self.config.num_steps)

    # Save a checkpoint of the final state
    sim.saveCheckpoint(str(self.output_dir / 'seg.chk'))

ContactMapRMSDReporter

Bases: OpenMMReporter

Reporter to compute contact maps and RMSD from an OpenMM simulation.

Source code in deepdrivewe/simulation/openmm.py 
 905
 906
 907
 908
 909
 910
 911
 912
 913
 914
 915
 916
 917
 918
 919
 920
 921
 922
 923
 924
 925
 926
 927
 928
 929
 930
 931
 932
 933
 934
 935
 936
 937
 938
 939
 940
 941
 942
 943
 944
 945
 946
 947
 948
 949
 950
 951
 952
 953
 954
 955
 956
 957
 958
 959
 960
 961
 962
 963
 964
 965
 966
 967
 968
 969
 970
 971
 972
 973
 974
 975
 976
 977
 978
 979
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
class ContactMapRMSDReporter(OpenMMReporter):
    """Reporter to compute contact maps and RMSD from an OpenMM simulation."""

    def __init__(
        self,
        report_interval: int,
        reference_file: Path,
        cutoff_angstrom: float = 8.0,
        mda_selection: str = 'protein and name CA',
        openmm_selection: Sequence[str] = ('CA',),
    ) -> None:
        """Initialize the reporter.

        Parameters
        ----------
        report_interval : int
            The interval at which to write frames.
        reference_file : Path
            The reference PDB file for the analysis.
        cutoff_angstrom : float
            The angstrom cutoff distance for defining contacts
            (default is 8.0).
        mda_selection : str
            The MDAnalysis selection string for the atoms to use
            (default is 'protein and name CA').
        openmm_selection : Sequence[str]
            The OpenMM selection strings for the atoms to use
            (default is ('CA',)).
        """
        super().__init__(report_interval)
        self.cutoff_angstrom = cutoff_angstrom
        self.openmm_selection = openmm_selection

        self._rows: list[np.ndarray] = []
        self._cols: list[np.ndarray] = []
        self._rmsd: list[float] = []

        # Load the reference structure and save the positions
        mda_u = MDAnalysis.Universe(reference_file)
        self._ref = mda_u.select_atoms(mda_selection).positions.copy()

    def get_contact_maps(self) -> np.ndarray:
        """Get the contact maps from the simulation.

        Returns
        -------
        np.ndarray
            The contact maps from the simulation as a ragged array
            shaped as (n_frames, *).
        """
        # Concatenate the row and col indices into a single array
        contact_maps = [
            np.concatenate(x)
            for x in zip(self._rows, self._cols, strict=False)
        ]

        # Collect the contact maps in a ragged numpy array
        contact_maps = np.array(contact_maps, dtype=object)

        return contact_maps

    def get_rmsds(self) -> np.ndarray:
        """Get the RMSDs from the simulation.

        Returns
        -------
        np.ndarray
            The RMSDs from the simulation shaped as (n_frames, 1).
        """
        return np.array(self._rmsd).reshape(-1, 1)

    def report(self, simulation: app.Simulation, state: openmm.State) -> None:
        """Generate a report.

        Parameters
        ----------
        simulation : Simulation
            The Simulation to generate a report for
        state : State
            The current state of the simulation
        """
        # Get the positions
        positions = self.get_positions(simulation, state)

        # Compute the contact map
        contact_map = distances.contact_matrix(
            positions,
            self.cutoff_angstrom,
            returntype='sparse',
        )

        # Convert the contact map to sparse format
        coo_matrix = contact_map.tocoo()

        # Append the row and col indices to lists
        self._rows.append(coo_matrix.row.astype('int16'))
        self._cols.append(coo_matrix.col.astype('int16'))

        # Compute the RMSD
        rmsd = rms.rmsd(positions, self._ref, superposition=True)
        self._rmsd.append(rmsd)

__init__ 

__init__(
    report_interval: int,
    reference_file: Path,
    cutoff_angstrom: float = 8.0,
    mda_selection: str = "protein and name CA",
    openmm_selection: Sequence[str] = ("CA",),
) -> None

Initialize the reporter.

Parameters:

Name Type Description Default
report_interval int

The interval at which to write frames.

 required
reference_file Path

The reference PDB file for the analysis.

 required
cutoff_angstrom float

The angstrom cutoff distance for defining contacts (default is 8.0).

 8.0
mda_selection str

The MDAnalysis selection string for the atoms to use (default is 'protein and name CA').

 'protein and name CA'
openmm_selection Sequence[str]

The OpenMM selection strings for the atoms to use (default is ('CA',)).

 ('CA',)
Source code in deepdrivewe/simulation/openmm.py 
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
def __init__(
    self,
    report_interval: int,
    reference_file: Path,
    cutoff_angstrom: float = 8.0,
    mda_selection: str = 'protein and name CA',
    openmm_selection: Sequence[str] = ('CA',),
) -> None:
    """Initialize the reporter.

    Parameters
    ----------
    report_interval : int
        The interval at which to write frames.
    reference_file : Path
        The reference PDB file for the analysis.
    cutoff_angstrom : float
        The angstrom cutoff distance for defining contacts
        (default is 8.0).
    mda_selection : str
        The MDAnalysis selection string for the atoms to use
        (default is 'protein and name CA').
    openmm_selection : Sequence[str]
        The OpenMM selection strings for the atoms to use
        (default is ('CA',)).
    """
    super().__init__(report_interval)
    self.cutoff_angstrom = cutoff_angstrom
    self.openmm_selection = openmm_selection

    self._rows: list[np.ndarray] = []
    self._cols: list[np.ndarray] = []
    self._rmsd: list[float] = []

    # Load the reference structure and save the positions
    mda_u = MDAnalysis.Universe(reference_file)
    self._ref = mda_u.select_atoms(mda_selection).positions.copy()

get_contact_maps 

get_contact_maps() -> np.ndarray

Get the contact maps from the simulation.

Returns:

Type Description
ndarray

The contact maps from the simulation as a ragged array shaped as (n_frames, *).
Source code in deepdrivewe/simulation/openmm.py 
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
def get_contact_maps(self) -> np.ndarray:
    """Get the contact maps from the simulation.

    Returns
    -------
    np.ndarray
        The contact maps from the simulation as a ragged array
        shaped as (n_frames, *).
    """
    # Concatenate the row and col indices into a single array
    contact_maps = [
        np.concatenate(x)
        for x in zip(self._rows, self._cols, strict=False)
    ]

    # Collect the contact maps in a ragged numpy array
    contact_maps = np.array(contact_maps, dtype=object)

    return contact_maps

get_rmsds 

get_rmsds() -> np.ndarray

Get the RMSDs from the simulation.

Returns:

Type Description
ndarray

The RMSDs from the simulation shaped as (n_frames, 1).
Source code in deepdrivewe/simulation/openmm.py 
966
967
968
969
970
971
972
973
974
def get_rmsds(self) -> np.ndarray:
    """Get the RMSDs from the simulation.

    Returns
    -------
    np.ndarray
        The RMSDs from the simulation shaped as (n_frames, 1).
    """
    return np.array(self._rmsd).reshape(-1, 1)

report 

report(simulation: Simulation, state: State) -> None

Generate a report.

Parameters:

Name Type Description Default
simulation Simulation

The Simulation to generate a report for

 required
state State

The current state of the simulation

 required
Source code in deepdrivewe/simulation/openmm.py 
 976
 977
 978
 979
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
def report(self, simulation: app.Simulation, state: openmm.State) -> None:
    """Generate a report.

    Parameters
    ----------
    simulation : Simulation
        The Simulation to generate a report for
    state : State
        The current state of the simulation
    """
    # Get the positions
    positions = self.get_positions(simulation, state)

    # Compute the contact map
    contact_map = distances.contact_matrix(
        positions,
        self.cutoff_angstrom,
        returntype='sparse',
    )

    # Convert the contact map to sparse format
    coo_matrix = contact_map.tocoo()

    # Append the row and col indices to lists
    self._rows.append(coo_matrix.row.astype('int16'))
    self._cols.append(coo_matrix.col.astype('int16'))

    # Compute the RMSD
    rmsd = rms.rmsd(positions, self._ref, superposition=True)
    self._rmsd.append(rmsd)