def listfamilies(element, with_description): """ Print on screen the list of installed PSF-pseudo families. """ from aiida import is_dbenv_loaded, load_dbenv if not is_dbenv_loaded(): load_dbenv() from aiida.orm import DataFactory from aiida_siesta.data.psf import PSFGROUP_TYPE PsfData = DataFactory('siesta.psf') from aiida.orm.querybuilder import QueryBuilder from aiida.orm.group import Group qb = QueryBuilder() qb.append(PsfData, tag='psfdata') if element: qb.add_filter(PsfData, {'attributes.element': {'in': element}}) qb.append( Group, group_of='psfdata', tag='group', project=["name", "description"], filters={ "type": { '==': PSFGROUP_TYPE } }) qb.distinct() if qb.count() > 0: for res in qb.dict(): group_name = res.get("group").get("name") group_desc = res.get("group").get("description") qb = QueryBuilder() qb.append( Group, tag='thisgroup', filters={ "name": { 'like': group_name } }) qb.append(PsfData, project=["id"], member_of='thisgroup') if with_description: description_string = ": {}".format(group_desc) else: description_string = "" click.echo("* {} [{} pseudos]{}".format(group_name, qb.count(), description_string)) else: click.echo("No valid PSF pseudopotential family found.", err=True)
def get_types(self): """ return available distinct types of nodes from database """ from aiida.orm.querybuilder import QueryBuilder qb = QueryBuilder() qb.append(self._aiida_class, project=["type"]) qb_response = qb.distinct().all() results = {} if len(qb_response) > 0: for ntype in qb_response: ntype = ntype[0] ntype_parts = ntype.split(".") if len(ntype_parts) > 0: dict_key = ntype_parts[0] if dict_key not in results.keys(): results[dict_key] = [] results[dict_key].append(ntype) for key, values in results.items(): results[key] = sorted(values) return results
def listfamilies(elements, with_description): """ Print on screen the list of upf families installed """ from aiida.orm import DataFactory from aiida.orm.data.upf import UPFGROUP_TYPE # pylint: disable=invalid-name UpfData = DataFactory('upf') from aiida.orm.querybuilder import QueryBuilder from aiida.orm.group import Group qb = QueryBuilder() qb.append(UpfData, tag='upfdata') if elements is not None: qb.add_filter(UpfData, {'attributes.element': {'in': elements}}) qb.append(Group, group_of='upfdata', tag='group', project=["name", "description"], filters={"type": { '==': UPFGROUP_TYPE }}) qb.distinct() if qb.count() > 0: for res in qb.dict(): group_name = res.get("group").get("name") group_desc = res.get("group").get("description") qb = QueryBuilder() qb.append(Group, tag='thisgroup', filters={"name": { 'like': group_name }}) qb.append(UpfData, project=["id"], member_of='thisgroup') if with_description: description_string = ": {}".format(group_desc) else: description_string = "" echo.echo_success("* {} [{} pseudos]{}".format( group_name, qb.count(), description_string)) else: echo.echo_warning("No valid UPF pseudopotential family found.")
def get_subtree(pk, level=0): qb = QueryBuilder() qb.append(cls=WorkCalculation, filters={'id': pk}, tag='workcalculation') qb.append( cls=WorkCalculation, project=['id'], # In the future, we should specify here the type of link # for now, CALL links are the only ones allowing calc-calc # (we here really want instead to follow CALL links) output_of='workcalculation', tag='subworkchains') result = list(itertools.chain(*qb.distinct().all())) # This will return a single flat list of tuples, where the first element # corresponds to the WorkChain pk and the second element is an integer # that represents its level of nesting within the chain return [(pk, level)] + list( itertools.chain( *[get_subtree(subpk, level=level + 1) for subpk in result]))
def query(datatype, project, past_days, group_pks, all_users): """ Perform the query """ import datetime from aiida.orm.implementation import Group from aiida.orm.user import User from aiida.orm.backend import construct_backend from aiida.orm.querybuilder import QueryBuilder from aiida.utils import timezone backend = construct_backend() qbl = QueryBuilder() if all_users is False: user = backend.users.get_automatic_user() qbl.append(User, tag="creator", filters={"email": user.email}) else: qbl.append(User, tag="creator") # If there is a time restriction data_filters = {} if past_days is not None: now = timezone.now() n_days_ago = now - datetime.timedelta(days=past_days) data_filters.update({"ctime": {'>=': n_days_ago}}) qbl.append(datatype, tag="data", created_by="creator", filters=data_filters, project=project) # If there is a group restriction if group_pks is not None: group_filters = dict() group_filters.update({"id": {"in": group_pks}}) qbl.append(Group, tag="group", filters=group_filters, group_of="data") qbl.order_by({datatype: {'ctime': 'asc'}}) object_list = qbl.distinct() return object_list.all()
def query_jobcalculations_by_computer_user_state( self, state, computer=None, user=None, only_computer_user_pairs=False, only_enabled=True, limit=None): """ Filter all calculations with a given state. Issue a warning if the state is not in the list of valid states. :param string state: The state to be used to filter (should be a string among those defined in aiida.common.datastructures.calc_states) :param computer: a Django DbComputer entry, or a Computer object, of a computer in the DbComputer table. A string for the hostname is also valid. :param user: a Django entry (or its pk) of a user in the DbUser table; if present, the results are restricted to calculations of that specific user :param bool only_computer_user_pairs: if False (default) return a queryset where each element is a suitable instance of Node (it should be an instance of Calculation, if everything goes right!) If True, return only a list of tuples, where each tuple is in the format ('dbcomputer__id', 'user__id') [where the IDs are the IDs of the respective tables] :param int limit: Limit the number of rows returned :return: a list of calculation objects matching the filters. """ # I assume that calc_states are strings. If this changes in the future, # update the filter below from dbattributes__tval to the correct field. from aiida.orm.computer import Computer from aiida.orm.calculation.job import JobCalculation from aiida.orm.user import User from aiida.orm.querybuilder import QueryBuilder from aiida.common.exceptions import InputValidationError from aiida.common.datastructures import calc_states if state not in calc_states: raise InputValidationError( "querying for calculation state='{}', but it " "is not a valid calculation state".format(state)) calcfilter = {'state': {'==': state}} computerfilter = {"enabled": {'==': True}} userfilter = {} if computer is None: pass elif isinstance(computer, int): computerfilter.update({'id': {'==': computer}}) elif isinstance(computer, Computer): computerfilter.update({'id': {'==': computer.pk}}) else: try: computerfilter.update({'id': {'==': computer.id}}) except AttributeError as e: raise Exception("{} is not a valid computer\n{}".format( computer, e)) if user is None: pass elif isinstance(user, int): userfilter.update({'id': {'==': user}}) else: try: userfilter.update({'id': {'==': int(user.id)}}) # Is that safe? except: raise Exception("{} is not a valid user".format(user)) qb = QueryBuilder() qb.append(type="computer", tag='computer', filters=computerfilter) qb.append(JobCalculation, filters=calcfilter, tag='calc', has_computer='computer') qb.append(type="user", tag='user', filters=userfilter, creator_of="calc") if only_computer_user_pairs: qb.add_projection("computer", "*") qb.add_projection("user", "*") returnresult = qb.distinct().all() else: qb.add_projection("calc", "*") if limit is not None: qb.limit(limit) returnresult = qb.all() returnresult = zip(*returnresult)[0] return returnresult
def get_bands_and_parents_structure(self, args): """ Search for bands and return bands and the closest structure that is a parent of the instance. This is the backend independent way, can be overriden for performance reason :returns: A list of sublists, each latter containing (in order): pk as string, formula as string, creation date, bandsdata-label """ import datetime from aiida.utils import timezone from aiida.orm.querybuilder import QueryBuilder from aiida.backends.utils import get_automatic_user from aiida.orm.implementation import User from aiida.orm.implementation import Group from aiida.orm.data.structure import (get_formula, get_symbols_string) from aiida.orm.data.array.bands import BandsData from aiida.orm.data.structure import StructureData qb = QueryBuilder() if args.all_users is False: au = get_automatic_user() user = User(dbuser=au) qb.append(User, tag="creator", filters={"email": user.email}) else: qb.append(User, tag="creator") bdata_filters = {} if args.past_days is not None: now = timezone.now() n_days_ago = now - datetime.timedelta(days=args.past_days) bdata_filters.update({"ctime": {'>=': n_days_ago}}) qb.append(BandsData, tag="bdata", created_by="creator", filters=bdata_filters, project=["id", "label", "ctime"]) group_filters = {} if args.group_name is not None: group_filters.update({"name": {"in": args.group_name}}) if args.group_pk is not None: group_filters.update({"id": {"in": args.group_pk}}) if group_filters: qb.append(Group, tag="group", filters=group_filters, group_of="bdata") qb.append( StructureData, tag="sdata", ancestor_of="bdata", # We don't care about the creator of StructureData project=["id", "attributes.kinds", "attributes.sites"]) qb.order_by({StructureData: {'ctime': 'desc'}}) list_data = qb.distinct() entry_list = [] already_visited_bdata = set() for [bid, blabel, bdate, sid, akinds, asites] in list_data.all(): # We process only one StructureData per BandsData. # We want to process the closest StructureData to # every BandsData. # We hope that the StructureData with the latest # creation time is the closest one. # This will be updated when the QueryBuilder supports # order_by by the distance of two nodes. if already_visited_bdata.__contains__(bid): continue already_visited_bdata.add(bid) if args.element is not None: all_symbols = [_["symbols"][0] for _ in akinds] if not any([s in args.element for s in all_symbols]): continue if args.element_only is not None: all_symbols = [_["symbols"][0] for _ in akinds] if not all([s in all_symbols for s in args.element_only]): continue # We want only the StructureData that have attributes if akinds is None or asites is None: continue symbol_dict = {} for k in akinds: symbols = k['symbols'] weights = k['weights'] symbol_dict[k['name']] = get_symbols_string(symbols, weights) try: symbol_list = [] for s in asites: symbol_list.append(symbol_dict[s['kind_name']]) formula = get_formula(symbol_list, mode=args.formulamode) # If for some reason there is no kind with the name # referenced by the site except KeyError: formula = "<<UNKNOWN>>" entry_list.append( [str(bid), str(formula), bdate.strftime('%d %b %Y'), blabel]) return entry_list