##############################################################################
# MDTraj: A Python Library for Loading, Saving, and Manipulating
# Molecular Dynamics Trajectories.
# Copyright 2012-2013 Stanford University and the Authors
#
# Authors: Robert McGibbon
# Contributors:
#
# MDTraj is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation, either version 2.1
# of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with MDTraj. If not, see <http://www.gnu.org/licenses/>.
##############################################################################
"""
Unit processing for MDTraj. This subpackage is a port of openmm.unit from
OpenMM. Unlike in openmm, the MDTraj library **does not pass around
"united quantities"**
The only publicly facing API from this package, for the purpose of MDTraj,
is "in_units_of", which does unit conversion of numbers or numpy arrays
where the input and output units are passed as strings.
"""
import ast
import sys
import numpy as np
from mdtraj.utils import import_
from mdtraj.utils.unit import unit_definitions
from mdtraj.utils.unit.quantity import Quantity
UNIT_DEFINITIONS = unit_definitions
try:
import openmm.unit as openmm_unit
except ImportError:
pass
__all__ = ["in_units_of"]
class _UnitContext(ast.NodeTransformer):
"""Node transformer for an AST hack that turns raw strings into
complex simt.unit.Unit expressions. See _str_to_unit for how this
is used -- it's not really meant to stand on its own
"""
# we want to do some validation to ensure that the AST only
# contains "safe" operations. These are operations that can reasonably
# appear in unit expressions
allowed_ops = [
ast.Expression,
ast.BinOp,
ast.Name,
ast.Attribute,
ast.Pow,
ast.Div,
ast.Mult,
ast.Constant,
]
def visit(self, node):
if not any(isinstance(node, a) for a in self.allowed_ops):
raise ValueError(
"Invalid unit expression. Contains dissallowed " "operation %s" % node.__class__.__name__,
)
return super().visit(node)
def visit_Name(self, node):
# we want to prefix all names to look like unit.nanometers instead
# of just "nanometers", because I don't want to import * from
# units into this module
if not hasattr(unit_definitions, node.id):
# also, let's take this opporunity to check that the node.id
# (which supposed to be the name of the unit, like "nanometers")
# is actually an attribute in openmm.unit
raise ValueError("%s is not a valid unit" % node.id)
return ast.Attribute(
value=ast.Name(id="unit_definitions", ctx=ast.Load()),
attr=node.id,
ctx=ast.Load(),
)
_unit_context = _UnitContext() # global instance of the visitor
def _str_to_unit(unit_string, openmm=False):
"""eval() based transformer that extracts a openmm.unit object
from a string description.
Parameters
----------
unit_string : str
string description of a unit. this may contain expressions with
multiplication, division, powers, etc.
Examples
--------
>>> type(_str_to_unit('nanometers**2/meters*gigajoules'))
<class 'openmm.unit.unit.Unit'>
>>> str(_str_to_unit('nanometers**2/meters*gigajoules'))
'nanometer**2*gigajoule/meter'
"""
# parse the string with the ast, and then run out unit context
# visitor on it, which will basically change bare names like
# "nanometers" into "unit.nanometers" and simulataniously check that
# there's no nefarious stuff in the expression.
assert isinstance(unit_string, str)
unit_definitions = UNIT_DEFINITIONS
if openmm:
unit_definitions = openmm_unit.unit_definitions
parsed = ast.parse(unit_string, mode="eval")
node = _unit_context.visit(parsed)
fixed_node = ast.fix_missing_locations(node)
output = eval(compile(fixed_node, "<string>", mode="eval"), {}, locals())
return output
[docs]
def in_units_of(quantity, units_in, units_out, inplace=False):
"""Convert a numerical quantity between unit systems.
Parameters
----------
quantity : {number, np.ndarray, openmm.unit.Quantity}
quantity can either be a unitted quantity -- i.e. instance of
openmm.unit.Quantity, or just a bare number or numpy array
units_in : str
If you supply a quantity that's not a openmm.unit.Quantity, you should
tell me what units it is in. If you don't, i'm just going to echo you
back your quantity without doing any unit checking.
units_out : str
A string description of the units you want out. This should look
like "nanometers/picosecond" or "nanometers**3" or whatever
inplace : bool
Attempt to do the transformation inplace, by mutating the `quantity`
argument and avoiding a copy. This is only possible if `quantity` is a
writable numpy array.
Returns
-------
rquantity : {number, np.ndarray}
The resulting quantity, in the new unit system. If the function was
called with `inplace=True` and `quantity` was a writable numpy array,
`rquantity` will alias the same memory as the input `quantity`, which
will have been changed inplace. Otherwise, if a copy was required,
`rquantity` will point to new memory.
Examples
--------
>>> in_units_of(1, 'meter**2/second', 'nanometers**2/picosecond')
1000000.0
"""
if quantity is None:
return quantity
if "openmm.unit" in sys.modules and isinstance(quantity, openmm_unit.Quantity):
units_in = quantity.unit
units_out = _str_to_unit(units_out, openmm=True)
quantity = quantity._value
elif isinstance(quantity, Quantity):
units_in = quantity.unit
units_out = _str_to_unit(units_out)
quantity = quantity._value
else:
if units_in is None:
return quantity
units_in = _str_to_unit(units_in)
units_out = _str_to_unit(units_out)
if not units_in.is_compatible(units_out):
raise TypeError(f'Unit "{units_in}" is not compatible with Unit "{units_out}".')
factor = units_in.conversion_factor_to(units_out)
if inplace and (isinstance(quantity, np.ndarray) and quantity.flags["WRITEABLE"]):
quantity *= factor
return quantity
return quantity * factor