Skip to content

Commit

Permalink
support for BandsData nodes without StructureData ancestors(#3817)
Browse files Browse the repository at this point in the history
When storing `BandsData` nodes without `StructureData` ancestors,  `verdi data bands list` would fail with a `KeyError` because it requires the formula from the connected structure.
This PR now uses formula `<<NOT FOUND>>` for isolated `BandsData` nodes and fixes `verdi data bands list` for this use case.
It also gets rid of a backend-specific implementation of `_extract_formula`.
  • Loading branch information
unkcpz authored Apr 11, 2020
1 parent 41a4e29 commit ff9e8a2
Show file tree
Hide file tree
Showing 3 changed files with 68 additions and 57 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(akinds, asites, args)
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
61 changes: 44 additions & 17 deletions aiida/backends/general/abstractqueries.py
Original file line number Diff line number Diff line change
Expand Up @@ -120,8 +120,19 @@ def get_statistics_dict(dataset):
return statistics

@staticmethod
def _extract_formula(args, akinds, asites):
"""Extract formula from the structure object."""
def _extract_formula(akinds, asites, args):
"""
Extract formula from the structure object.
:param akinds: list of kinds, e.g. [{'mass': 55.845, 'name': 'Fe', 'symbols': ['Fe'], 'weights': [1.0]},
{'mass': 15.9994, 'name': 'O', 'symbols': ['O'], 'weights': [1.0]}]
:param asites: list of structure sites e.g. [{'position': [0.0, 0.0, 0.0], 'kind_name': 'Fe'},
{'position': [2.0, 2.0, 2.0], 'kind_name': 'O'}]
:param args: a namespace with parsed command line parameters, here only 'element' and 'element_only' are used
:type args: dict
:return: a string with formula if the formula is found
"""
from aiida.orm.nodes.data.structure import (get_formula, get_symbols_string)

if args.element is not None:
Expand All @@ -136,7 +147,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 +172,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 +186,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 +214,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 +234,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(akinds, asites, args)
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 ff9e8a2

Please sign in to comment.