Source code for

A collection of classes used to query a storage backend for
data which matches a set of criteria.
import abc

from openff.evaluator.attributes import UNDEFINED, Attribute, AttributeClass
from openff.evaluator.datasets import PropertyPhase
from openff.evaluator.forcefield import ForceFieldSource
from import QueryAttribute
from import (
from openff.evaluator.substances import ExactAmount, Substance
from openff.evaluator.thermodynamics import ThermodynamicState

[docs]class BaseDataQuery(AttributeClass, abc.ABC): """A base class for queries which can be made to a `StorageBackend`. """
[docs] @classmethod @abc.abstractmethod def data_class(cls): """The type of data class that this query can be applied to. Returns ------- type of BaseStoredData """ raise NotImplementedError()
[docs] def apply(self, data_object): """Apply this query to a data object. Parameters ---------- data_object: BaseStoredData The data object to apply the query to. Returns ------- tuple of Any, optional The values of the matched parameters of the data object fully matched this query, otherwise `None`. """ if not isinstance(data_object, self.data_class()): return None matches = [] for attribute_name in self.get_attributes(QueryAttribute): attribute = getattr(self.__class__, attribute_name) if not hasattr(data_object, attribute_name) or attribute.custom_match: continue query_value = getattr(self, attribute_name) if query_value == UNDEFINED: continue data_value = getattr(data_object, attribute_name) matches.append(None if data_value != query_value else data_value) if any(x is None for x in matches): return None return tuple(matches)
[docs] @classmethod def from_data_object(cls, data_object): """Returns the query which would match this data object. Parameters ---------- data_object: BaseStoredData The data object to construct the query for. Returns ------- cls The query which would match this data object. """ query = cls() for attribute_name in cls.get_attributes(): if not hasattr(data_object, attribute_name): continue attribute_value = getattr(data_object, attribute_name) setattr(query, attribute_name, attribute_value) return query
[docs]class SubstanceQuery(AttributeClass, abc.ABC): """A query which focuses on finding data which was collected for substances with specific traits, e.g which contains both a solute and solvent, or only a solvent etc. """ components_only = Attribute( docstring="Only match pure data which was collected for " "one of the components in the query substance.", type_hint=bool, default_value=False, ) # component_roles = QueryAttribute( # docstring="Returns data for only the subset of a substance " # "which has the requested roles.", # type_hint=list, # optional=True, # )
[docs] def validate(self, attribute_type=None): super(SubstanceQuery, self).validate(attribute_type)
# if ( # self.components_only # and self.component_roles != UNDEFINED # and len(self.components_only) > 0 # ): # # raise ValueError( # "The `component_roles` attribute cannot be used when " # "the `components_only` attribute is `True`." # )
[docs]class ForceFieldQuery(BaseDataQuery): """A class used to query a `StorageBackend` for `ForceFieldData` which meet the specified criteria. """
[docs] @classmethod def data_class(cls): return ForceFieldData
force_field_source = QueryAttribute( docstring="The force field source to query for.", type_hint=ForceFieldSource, optional=True, )
[docs]class BaseSimulationDataQuery(BaseDataQuery, abc.ABC): """The base class for queries which will retrieve ``BaseSimulationData`` derived data. """ substance = QueryAttribute( docstring="The substance which the data should have been collected " "for. Data for a subset of this substance can be queried for by " "using the `substance_query` attribute", type_hint=Substance, optional=True, custom_match=True, ) substance_query = QueryAttribute( docstring="The subset of the `substance` to query for. This option " "can only be used when the `substance` attribute is set.", type_hint=SubstanceQuery, optional=True, custom_match=True, ) thermodynamic_state = QueryAttribute( docstring="The state at which the data should have been collected.", type_hint=ThermodynamicState, optional=True, ) property_phase = QueryAttribute( docstring="The phase of the substance (e.g. liquid, gas).", type_hint=PropertyPhase, optional=True, ) source_calculation_id = QueryAttribute( docstring="The server id which should have generated this data.", type_hint=str, optional=True, ) force_field_id = QueryAttribute( docstring="The id of the force field parameters which used to " "generate the data.", type_hint=str, optional=True, ) def _match_substance(self, data_object): """Attempt to match the substance (or a subset of it). Parameters ---------- data_object: StoredSimulationData The data object to match against. Returns ------- Substance, optional The matched substance if a match is made, otherwise `None`. """ if self.substance == UNDEFINED: return None data_substance: Substance = data_object.substance if self.substance_query == UNDEFINED: return None if self.substance != data_substance else self.substance # Handle the sub-substance match. if self.substance_query.components_only: if data_substance.number_of_components != 1: # We are only interested in pure data. return None for component in self.substance.components: if component.smiles != data_substance.components[0].smiles: continue # Make sure the amount type matches up i.e either both # are defined in mole fraction, or both as an exact amount. data_amount = next( iter(data_substance.get_amounts(component.identifier)) ) query_amount = next( iter(self.substance.get_amounts(component.identifier)) ) if type(data_amount) != type(query_amount): continue if isinstance(data_amount, ExactAmount) and data_amount != query_amount: # Make sure there is the same amount if we are # dealing with exact amounts. continue # A match was found. return data_substance return None
[docs] def apply(self, data_object, attributes_to_ignore=None): matches = [] # Apply a custom match behaviour for the substance # attribute. if self.substance != UNDEFINED: matches.append(self._match_substance(data_object)) base_matches = super(BaseSimulationDataQuery, self).apply(data_object) base_matches = [None] if base_matches is None else base_matches matches = [*matches, *base_matches] if len(matches) == 0 or any(x is None for x in matches): return None return tuple(matches)
[docs] def validate(self, attribute_type=None): super(BaseSimulationDataQuery, self).validate(attribute_type) if self.substance_query != UNDEFINED and self.substance == UNDEFINED: raise ValueError( "The `substance_query` can only be used when the " "`substance` attribute is set." )
[docs]class SimulationDataQuery(BaseSimulationDataQuery): """A class used to query a ``StorageBackend`` for ``StoredSimulationData`` objects which meet the specified set of criteria. """
[docs] @classmethod def data_class(cls): return StoredSimulationData
number_of_molecules = QueryAttribute( docstring="The total number of molecules in the system.", type_hint=int, optional=True, )
[docs]class FreeEnergyDataQuery(BaseSimulationDataQuery): """A class used to query a ``StorageBackend`` for ``FreeEnergyData`` objects which meet the specified set of criteria. """
[docs] @classmethod def data_class(cls): return StoredFreeEnergyData