python_packaging/src/gmxapi/simulation/read_tpr.py

   1 #
   2 # This file is part of the GROMACS molecular simulation package.
   3 #
   4 # Copyright (c) 2019, by the GROMACS development team, led by
   5 # Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
   6 # and including many others, as listed in the AUTHORS file in the
   7 # top-level source directory and at http://www.gromacs.org.
   8 #
   9 # GROMACS is free software; you can redistribute it and/or
  10 # modify it under the terms of the GNU Lesser General Public License
  11 # as published by the Free Software Foundation; either version 2.1
  12 # of the License, or (at your option) any later version.
  13 #
  14 # GROMACS is distributed in the hope that it will be useful,
  15 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  16 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  17 # Lesser General Public License for more details.
  18 #
  19 # You should have received a copy of the GNU Lesser General Public
  20 # License along with GROMACS; if not, see
  21 # http://www.gnu.org/licenses, or write to the Free Software Foundation,
  22 # Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
  23 #
  24 # If you want to redistribute modifications to GROMACS, please
  25 # consider that scientific software is very special. Version
  26 # control is crucial - bugs must be traceable. We will be happy to
  27 # consider code for inclusion in the official distribution, but
  28 # derived work must not be called official GROMACS. Details are found
  29 # in the README & COPYING files - if they are missing, get the
  30 # official version at http://www.gromacs.org.
  31 #
  32 # To help us fund GROMACS development, we humbly ask that you cite
  33 # the research papers on the package. Check out http://www.gromacs.org.
  34
  35 """read_tpr operation module
  36
  37 Provides implementation classes and user interface for gmxapi.read_tpr.
  38 """
  39
  40 import inspect
  41 import typing
  42
  43 import gmxapi
  44 import gmxapi.abc
  45 import gmxapi.exceptions
  46 import gmxapi.operation as _op
  47 import gmxapi.typing
  48 from gmxapi.operation import InputCollectionDescription
  49 from gmxapi.simulation.abc import ModuleObject
  50
  51 from . import fileio
  52
  53 # Initialize module-level logger
  54 from gmxapi import logger as root_logger
  55
  56 logger = root_logger.getChild('read_tpr')
  57 logger.info('Importing {}'.format(__name__))
  58
  59
  60 #
  61 # Interface classes and internal details
  62 #
  63
  64 # TODO: The output of read_tpr and modify_input should be the same and part of the
  65 #  simulation module specification. Such output is either a special type of output proxy
  66 #  or Future.
  67 _output_descriptors = (
  68     _op.OutputDataDescriptor('_simulation_input', str),
  69     _op.OutputDataDescriptor('parameters', dict)
  70 )
  71 _publishing_descriptors = {desc._name: gmxapi.operation.Publisher(desc._name, desc._dtype) for desc in
  72                            _output_descriptors}
  73 _output = _op.OutputCollectionDescription(**{descriptor._name: descriptor._dtype for descriptor in
  74                                              _output_descriptors})
  75
  76
  77
  78 class OutputDataProxy(ModuleObject,
  79                       _op.DataProxyBase,
  80                       descriptors=_output_descriptors):
  81     """Implement the 'output' attribute of ReadTpr operations."""
  82     def __init__(self, *args, **kwargs):
  83         _op.DataProxyBase.__init__(self, *args, **kwargs)
  84
  85
  86 class PublishingDataProxy(_op.DataProxyBase,
  87                           descriptors=_publishing_descriptors
  88                           ):
  89     """Manage output resource updates for ReadTpr operation."""
  90
  91
  92 _output_factory = _op.OutputFactory(output_proxy=OutputDataProxy,
  93                                     output_description=_output,
  94                                     publishing_data_proxy=PublishingDataProxy)
  95
  96
  97 class SessionResources(object):
  98     """Input and output run-time resources for a ReadTpr operation."""
  99     def __init__(self, tpr_filename, publisher: PublishingDataProxy):
 100         self.tpr_object = fileio.TprFile(filename=tpr_filename, mode='r')
 101         self.output = publisher
 102
 103
 104 #
 105 # Helpers
 106 #
 107
 108
 109 _input = _op.InputCollectionDescription(
 110     [('filename', inspect.Parameter('filename',
 111                                     inspect.Parameter.POSITIONAL_OR_KEYWORD,
 112                                     annotation=str))])
 113
 114
 115 # TODO: Clarify. The actual input and output arguments passed are customized for this operation.
 116 def _session_resource_factory(input: _op.InputPack, output: 'PublishingDataProxy') -> SessionResources:
 117     """Translate resources from the gmxapi.operation Context to the ReadTpr implementation."""
 118     filename = input.kwargs['filename']
 119     return SessionResources(tpr_filename=filename, publisher=output)
 120
 121
 122 def _standard_node_resource_factory(*args, **kwargs):
 123     """Translate Python UI input to the gmxapi.operation node builder inputs."""
 124     return _input.bind(*args, **kwargs)
 125
 126
 127 def _run(resources: SessionResources):
 128     """Operation implementation in the gmxapi.operation module context."""
 129     # TODO: Implement for other source/target contexts. We don't always need to
 130     #  produce all outputs.
 131     with resources.tpr_object as fh:
 132         params = fh._tprFileHandle.params().extract()
 133         resources.output.parameters = params
 134         resources.output._simulation_input = fh.filename
 135
 136
 137 #
 138 # Implementation
 139 #
 140
 141 # Note: we borrow the implementation from operation.ResourceManager for now,
 142 # but in the future we want the implementations to either be decoupled or
 143 # for implementations in a given context to be coupled to details that are clearly
 144 # and explicitly related to that context. Right now, operation.ResourceManager
 145 # is tied to the implementation of Contexts in gmxapi.operation, but that is not
 146 # sufficiently clear and explicit.
 147 class ResourceManager(gmxapi.operation.ResourceManager):
 148     """Manage resources for the ReadTpr operation in the gmxapi.operation contexts.
 149
 150     Extends gmxapi.operation.ResourceManager to tolerate non-standard data payloads.
 151     Futures managed by this resource manager may contain additional attributes.
 152     """
 153     def future(self, name: str, description: _op.ResultDescription):
 154         tpr_future = super().future(name=name, description=description)
 155         return tpr_future
 156
 157     def data(self) -> OutputDataProxy:
 158         return OutputDataProxy(self)
 159
 160     def update_output(self):
 161         logger.debug('Updating output for {}.'.format(self.operation_id))
 162         super().update_output()
 163         for descriptor in _output_descriptors:
 164             name = descriptor._name
 165             if not self.is_done(name):
 166                 raise gmxapi.exceptions.ApiError('Expected output {} not updated.'.format(name))
 167
 168 # TODO: Consider making Generic in source and target context type variables,
 169 #  or leave unspecified and use generic function or pair of single_dispatch functions.
 170 #  Need to know the right home for input_description, if needed.
 171 class ResourceFactory(gmxapi.abc.ResourceFactory):
 172     """ReadTpr resource factory.
 173
 174     Generic class for creating resources passed to read_tpr implementation details.
 175     Dispatching may occur based on the source and target Context of factory action.
 176     """
 177     def __init__(self, target_context, source_context):
 178         """Initialize an instance to support read_tpr action.
 179
 180         Arguments:
 181             source_context: The source of the resources in the calling scope.
 182             target_context: The Context in which the factory product will be consumed.
 183
 184         """
 185         self.source_context = source_context  # Determine input form of *create* method.
 186         self.target_context = target_context  # Context in which resources are consumed.
 187
 188     # TODO: clean up the dispatching. What to do when input comes from multiple sources?
 189     # Use a typing overload or a single-dispatch functor for clearer input/result typing.
 190     def __call__(self, *args, **kwargs):
 191         """Create the resource product.
 192
 193         Dispatch to appropriate factory functionality based on the *source_context*
 194         and *target_context*.
 195         """
 196         # context is used for dispatching and annotates the Context of the other arguments.
 197         # context == None implies inputs are from Python UI.
 198         if self.source_context is None:
 199             if isinstance(self.target_context, _op.Context):
 200                 return _standard_node_resource_factory(*args, **kwargs)
 201         if isinstance(self.source_context, _op.Context):
 202             # TODO: Check whether the consumer is a Context.NodeBuilder or an operation runner.
 203             # We don't yet use this dispatcher for building nodes, so assume we are launching a session.
 204             assert 'input' in kwargs
 205             assert 'output' in kwargs
 206             return _session_resource_factory(input=kwargs['input'], output=kwargs['output'])
 207         raise gmxapi.exceptions.NotImplementedError(
 208             'No translation from {} context to {}'.format(self.source_context, self.target_context))
 209
 210     @typing.overload
 211     def input_description(self, context: _op.Context) -> _op.InputDescription:
 212         ...
 213
 214     def input_description(self, context: gmxapi.abc.Context):
 215         """Get an input description usable in the indicated Context.
 216
 217         Arguments:
 218             context: Context for which to dispatch the generation of input description.
 219
 220         Overrides gmxapi.abc.ResourceFactory for collaboration between this
 221         resource factory and a target Context.
 222         """
 223         # TODO: Clarify exposure and scope as gmxapi.abc.ResourceFactory is refined.
 224         # The expected use case is that a Context implementation may consult the
 225         # input_description when preparing input to a resource factory, or even
 226         # when determining in which Context a resource should be created or an
 227         # operation dispatched. It seems sensible that the ResourceFactory should
 228         # be able to describe its possible inputs. Instance functions are the
 229         # most common and simple Python detail, so it makes sense for this to be
 230         # an instance method instead of classmethod or staticmethod, until shown
 231         # otherwise.
 232         # Also note that ResourceFactory is likely to be a Generic class or a
 233         # container for composed functionality. Customizing method signatures
 234         # based on instance data is more brittle than behavior determined at the
 235         # class level. The behavior of this function is determined at the class
 236         # level to be a dispatcher, which may utilize instance data, but as an
 237         # implementation detail that is not the business of the caller.
 238         # This should probably be implemented in terms of the standard Python
 239         # functools.single_dispatch generic function, but it is cleaner if the
 240         # method itself is not generic beyond the level of typing overloads.
 241         if isinstance(context, _op.Context):
 242             return StandardInputDescription()
 243         raise gmxapi.exceptions.NotImplementedError('No input description available for {} context'.format(context))
 244
 245
 246 class StandardInputDescription(_op.InputDescription):
 247     """Provide the ReadTpr input description in gmxapi.operation Contexts."""
 248
 249     # TODO: Improve fingerprinting.
 250     # If _make_uid can't make sufficiently unique IDs, use additional "salt":
 251     # _next_uid = 0
 252
 253     @staticmethod
 254     def _make_uid(input) -> str:
 255         # TODO: Use input fingerprint for more useful identification.
 256         salt = hash(input)
 257         # If can't make sufficiently unique IDs, use additional "salt"
 258         # from a class data member. E.g.
 259         #     new_uid = 'read_tpr_{}_{}'.format(_next_uid, salt)
 260         #     _next_uid += 1
 261         new_uid = 'read_tpr_{}'.format(salt)
 262         return new_uid
 263
 264     def signature(self) -> InputCollectionDescription:
 265         return _input
 266
 267     def make_uid(self, input: _op.DataEdge) -> str:
 268         assert isinstance(input, _op.DataEdge)
 269         return self._make_uid(input)
 270
 271
 272 class RegisteredOperation(_op.OperationImplementation, metaclass=_op.OperationMeta):
 273     """Provide the gmxapi compatible ReadTpr implementation."""
 274
 275     # This is a class method to allow the class object to be used in gmxapi.operation._make_registry_key
 276     @classmethod
 277     def name(self) -> str:
 278         """Canonical name for the operation."""
 279         return 'read_tpr'
 280
 281     @classmethod
 282     def namespace(self) -> str:
 283         """read_tpr is importable from the gmxapi module."""
 284         return 'gmxapi'
 285
 286     @classmethod
 287     def director(cls, context: gmxapi.abc.Context) -> _op.OperationDirector:
 288         if isinstance(context, _op.Context):
 289             return StandardDirector(context)
 290         raise gmxapi.exceptions.NotImplementedError(
 291             'No dispatcher for context {} of type {}'.format(context, type(context)))
 292
 293
 294 class StandardOperationHandle(_op.AbstractOperation, ModuleObject):
 295     """Handle used in Python UI or gmxapi.operation Contexts."""
 296     def __init__(self, resource_manager: ResourceManager):
 297         self.__resource_manager = resource_manager
 298
 299     def run(self):
 300         self.__resource_manager.update_output()
 301
 302     @property
 303     def output(self) -> OutputDataProxy:
 304         return self.__resource_manager.data()
 305
 306
 307 class StandardDirector(gmxapi.abc.OperationDirector):
 308     """Direct the instantiation of a read_tpr node in a gmxapi.operation Context.
 309
 310     .. todo:: Compose this behavior in a more generic class.
 311
 312     .. todo:: Describe where instances live.
 313     """
 314     def __init__(self, context: _op.Context):
 315         if not isinstance(context, _op.Context):
 316             raise gmxapi.exceptions.ValueError('StandardDirector requires a gmxapi.operation Context.')
 317         self.context = context
 318
 319     def __call__(self, resources: _op.DataSourceCollection, label: str) -> StandardOperationHandle:
 320         builder = self.context.node_builder(operation=RegisteredOperation, label=label)
 321
 322         builder.set_resource_factory(_session_resource_factory)
 323         builder.set_input_description(StandardInputDescription())
 324         builder.set_handle(StandardOperationHandle)
 325
 326         def runner_director(resources):
 327             def runner():
 328                 _run(resources)
 329             return runner
 330
 331         builder.set_runner_director(runner_director)
 332         builder.set_output_factory(_output_factory)
 333
 334         # Note: we have not yet done any dispatching based on *resources*. We should
 335         # translate the resources provided into the form that the Context can subscribe to
 336         # using the dispatching resource_factory. In the second draft, this operation
 337         # is dispatched to a SimulationModuleContext, which can be subscribed directly
 338         # to sources that are either literal filenames or gmxapi.simulation sources,
 339         # while standard Futures can be resolved in the standard context.
 340         #
 341         assert isinstance(resources, _op.DataSourceCollection)
 342         assert 'filename' in resources
 343         builder.add_input('filename', resources['filename'])
 344
 345         handle = builder.build()
 346         assert isinstance(handle, StandardOperationHandle)
 347         return handle
 348
 349     def handle_type(self, context: gmxapi.abc.Context):
 350         return StandardOperationHandle
 351
 352     def resource_factory(self,
 353                          source: typing.Union[gmxapi.abc.Context, None],
 354                          target: gmxapi.abc.Context = None) -> ResourceFactory:
 355         # Distinguish between the UIContext, in which input is in the form
 356         # of function call arguments, and the StandardContext, implemented in
 357         # gmxapi.operation. UIContext is probably a virtual context that is
 358         # asserted by callers in order to get a factory that normalizes UI input
 359         # for the StandardContext.
 360         #
 361         if target is None:
 362             target = self.context
 363         if source is None:
 364             if isinstance(target, _op.Context):
 365                 return ResourceFactory(target_context=target, source_context=source)
 366         if isinstance(source, _op.Context):
 367             return ResourceFactory(target_context=target, source_context=source)
 368         raise gmxapi.exceptions.ValueError('No dispatching from {} context to {}'.format(source, target))
 369
 370
 371 def read_tpr(filename, label: str = None, context=None):
 372     """
 373
 374     Arguments:
 375         filename: input file name
 376         label: optional human-readable label with which to tag the new node
 377         context: Context in which to return a handle to the new node.
 378                  Use default (None) for Python scripting interface
 379
 380     Returns:
 381         Reference (handle) to the new operation instance (node).
 382
 383     """
 384     handle_context = context
 385     if handle_context is not None:
 386         raise gmxapi.exceptions.NotImplementedError(
 387             'context must be None. This factory is only for the Python UI right now.')
 388
 389     # 1. Handle node creation in the scripting interface.
 390     # When *context* input is None, dispatch to the current Context. Confirm that
 391     # it is a standard context from the gmxapi.operation module, translate the
 392     # input into that context, and create the node.
 393
 394     # 2. Dispatch to SimulationModuleContext.
 395     # Operation is not fully implemented in gmxapi.operation context. When creating
 396     # a node in a gmxapi.operation context, dispatch and subscribe to a SimulationModuleContext,
 397     # in which
 398
 399     # 3. Handle consumption in SimulationModuleContext.
 400     # When a consuming operation is native to the SimulationModuleContext,
 401     # detect that a richer interface is available and use it. Possible implementation:
 402     # Chain of Responsibility: subscribe() is serviced by the Context "closest"
 403     # to the source. subscriber is the most native compatible Context for the consuming operation
 404     # but decays to the context in which the handle is being created. subscribe()
 405     # can accept a callback function and an indication of the form of the message
 406     # to send (a consuming Context or ResourceFactory).
 407
 408     # TODO: Other types of input, such as File placeholders.
 409
 410     target_context = _op.current_context()
 411     assert isinstance(target_context, _op.Context)
 412     # Get a director that will create a node in the standard context.
 413     node_director = _op._get_operation_director(RegisteredOperation, context=target_context)
 414     assert isinstance(node_director, StandardDirector)
 415     # TODO: refine this protocol
 416     assert handle_context is None
 417     resource_factory = node_director.resource_factory(source=handle_context, target=target_context)
 418     resources = resource_factory(filename=filename)
 419     handle = node_director(resources=resources, label=label)
 420     # Note: One effect of the assertions above is to help the type checker infer
 421     # the return type of the handle. It is hard to convince the type checker that
 422     # the return value of the node builder is up-cast. We might be able to resolve
 423     # this by making both get_operation_director and ReadTprImplementation
 424     # generics of the handle type, or using the handle type as the key for
 425     # get_operation_director.
 426     return handle