Skip to content

Commit

Permalink
isolated bandsdata list as 'NOT FOUND' and some test
Browse files Browse the repository at this point in the history
both backend supported
  • Loading branch information
unkcpz committed Mar 3, 2020
1 parent d44dcbd commit 490a8f4
Show file tree
Hide file tree
Showing 3 changed files with 55 additions and 55 deletions.
51 changes: 12 additions & 39 deletions aiida/backends/djsite/queries.py
Original file line number Diff line number Diff line change
Expand Up @@ -108,43 +108,6 @@ def query_group(q_object, args):
if args.group_pk is not None:
q_object.add(Q(dbgroups__pk__in=args.group_pk), Q.AND)

@staticmethod
def _extract_formula(struc_pk, args, deser_data):
"""Extract formula."""
from aiida.orm.nodes.data.structure import (get_formula, get_symbols_string)

if struc_pk is not None:
# Exclude structures by the elements
if args.element is not None:
all_kinds = [k['symbols'] for k in deser_data[struc_pk]['kinds']]
all_symbols = [item for sublist in all_kinds for item in sublist]
if not any([s in args.element for s in all_symbols]):
return None
if args.element_only is not None:
all_kinds = [k['symbols'] for k in deser_data[struc_pk]['kinds']]
all_symbols = [item for sublist in all_kinds for item in sublist]
if not all([s in all_symbols for s in args.element_only]):
return None

# build the formula
symbol_dict = {
k['name']: get_symbols_string(k['symbols'], k['weights']) for k in deser_data[struc_pk]['kinds']
}
try:
symbol_list = [symbol_dict[s['kind_name']] for s in deser_data[struc_pk]['sites']]
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>>'
# cycle if we imposed the filter on elements
if args.element is not None or args.element_only is not None:
return None
else:
formula = '<<UNKNOWN>>'

return formula

def get_bands_and_parents_structure(self, args):
"""Returns bands and closest parent structure."""
from django.db.models import Q
Expand Down Expand Up @@ -175,14 +138,24 @@ def get_bands_and_parents_structure(self, args):
# get the closest structures (WITHOUT DbPath)
structure_dict = get_closest_parents(pks, Q(node_type='data.structure.StructureData.'), chunk_size=1)

struc_pks = [structure_dict[pk] for pk in pks]
struc_pks = [structure_dict.get(pk) for pk in pks]

# query for the attributes needed for the structure formula
res_attr = models.DbNode.objects.filter(id__in=struc_pks).values_list('id', 'attributes')
res_attr = {rattr[0]: rattr[1] for rattr in res_attr}

# prepare the printout
for (b_id_lbl_date, struc_pk) in zip(this_chunk, struc_pks):
formula = self._extract_formula(struc_pk, args, {rattr[0]: rattr[1] for rattr in res_attr})
if struc_pk is not None:
strct = res_attr[struc_pk]
akinds, asites = strct['kinds'], strct['sites']
formula = self._extract_formula(args, akinds, asites)
else:
if args.element is not None or args.element_only is not None:
formula = None
else:
formula = 'NOT FOUND'

if formula is None:
continue
entry_list.append([
Expand Down
46 changes: 31 additions & 15 deletions aiida/backends/general/abstractqueries.py
Original file line number Diff line number Diff line change
Expand Up @@ -136,7 +136,7 @@ def _extract_formula(args, akinds, asites):

# We want only the StructureData that have attributes
if akinds is None or asites is None:
return None
return '<<UNKNOWN>>'

symbol_dict = {}
for k in akinds:
Expand All @@ -161,7 +161,9 @@ def get_bands_and_parents_structure(self, args):
:returns:
A list of sublists, each latter containing (in order):
pk as string, formula as string, creation date, bandsdata-label"""
pk as string, formula as string, creation date, bandsdata-label
"""
# pylint: disable=too-many-locals

import datetime
from aiida.common import timezone
Expand All @@ -173,22 +175,23 @@ def get_bands_and_parents_structure(self, args):
else:
q_build.append(orm.User, tag='creator')

bdata_filters = {}
if args.past_days is not None:
bdata_filters.update({'ctime': {'>=': timezone.now() - datetime.timedelta(days=args.past_days)}})

q_build.append(
orm.BandsData, tag='bdata', with_user='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:
q_build.append(orm.Group, tag='group', filters=group_filters, with_node='bdata')

q_build.append(orm.Group, tag='group', filters=group_filters, with_user='creator')

bdata_filters = {}
if args.past_days is not None:
bdata_filters.update({'ctime': {'>=': timezone.now() - datetime.timedelta(days=args.past_days)}})

q_build.append(
orm.BandsData, tag='bdata', with_group='group', filters=bdata_filters, project=['id', 'label', 'ctime']
)
bands_list_data = q_build.all()

q_build.append(
orm.StructureData,
Expand All @@ -200,12 +203,15 @@ def get_bands_and_parents_structure(self, args):

q_build.order_by({orm.StructureData: {'ctime': 'desc'}})

list_data = q_build.distinct()
structure_dict = dict()
list_data = q_build.distinct().all()
for bid, _, _, _, akinds, asites in list_data:
structure_dict[bid] = (akinds, asites)

entry_list = []
already_visited_bdata = set()

for [bid, blabel, bdate, _, akinds, asites] in list_data.all():
for [bid, blabel, bdate] in bands_list_data:

# We process only one StructureData per BandsData.
# We want to process the closest StructureData to
Expand All @@ -217,7 +223,17 @@ def get_bands_and_parents_structure(self, args):
if already_visited_bdata.__contains__(bid):
continue
already_visited_bdata.add(bid)
formula = self._extract_formula(args, akinds, asites)
strct = structure_dict.get(bid, None)

if strct is not None:
akinds, asites = strct
formula = self._extract_formula(args, akinds, asites)
else:
if args.element is not None or args.element_only is not None:
formula = None
else:
formula = 'NOT FOUND'

if formula is None:
continue
entry_list.append([str(bid), str(formula), bdate.strftime('%d %b %Y'), blabel])
Expand Down
13 changes: 12 additions & 1 deletion tests/cmdline/commands/test_data.py
Original file line number Diff line number Diff line change
Expand Up @@ -141,7 +141,6 @@ def data_listing_test(self, datatype, search_string, ids):

# Check that the past days filter works as expected
past_days_flags = ['-p', '--past-days']
# past_days_flags = ['-p']
for flag in past_days_flags:
options = [flag, '1']
res = self.cli_runner.invoke(listing_cmd, options, catch_exceptions=False)
Expand All @@ -158,6 +157,7 @@ def data_listing_test(self, datatype, search_string, ids):
)

# Check that the group filter works as expected
# if ids is not None:
group_flags = ['-G', '--groups']
for flag in group_flags:
# Non empty group
Expand Down Expand Up @@ -289,10 +289,14 @@ def connect_structure_bands(strct): # pylint: disable=unused-argument

bands = connect_structure_bands(strct)

bands_isolated = BandsData()
bands_isolated.store()

# Create 2 groups and add the data to one of them
g_ne = Group(label='non_empty_group')
g_ne.store()
g_ne.add_nodes(bands)
g_ne.add_nodes(bands_isolated)

g_e = Group(label='empty_group')
g_e.store()
Expand Down Expand Up @@ -321,6 +325,13 @@ def test_bandlistshelp(self):

def test_bandslist(self):
self.data_listing_test(BandsData, 'FeO', self.ids)
self.data_listing_test(BandsData, 'NOT FOUND', self.ids)

def test_bandslist_with_elements(self):
options = ['-e', 'Fe']
res = self.cli_runner.invoke(cmd_bands.bands_list, options, catch_exceptions=False)
self.assertIn(b'FeO', res.stdout_bytes, 'The string "FeO" was not found in the listing')
self.assertNotIn(b'NOT FOUND', res.stdout_bytes, 'The string "NOT FOUND" should not in the listing')

def test_bandexporthelp(self):
output = sp.check_output(['verdi', 'data', 'bands', 'export', '--help'])
Expand Down

0 comments on commit 490a8f4

Please sign in to comment.