def test_with_subclasses(self, computer):
extra_name = self.__class__.__name__ + "/test_with_subclasses"
calc_params = {
'computer': computer,
'resources': {
'num_machines': 1,
'num_mpiprocs_per_machine': 1
}
}
TemplateReplacerCalc = CalculationFactory(
'simpleplugins.templatereplacer')
ParameterData = DataFactory('parameter')
a1 = JobCalculation(**calc_params).store()
# To query only these nodes later
a1.set_extra(extra_name, True)
a2 = TemplateReplacerCalc(**calc_params).store()
# To query only these nodes later
a2.set_extra(extra_name, True)
a3 = Data().store()
a3.set_extra(extra_name, True)
a4 = ParameterData(dict={'a': 'b'}).store()
a4.set_extra(extra_name, True)
a5 = Node().store()
a5.set_extra(extra_name, True)
# I don't set the extras, just to be sure that the filtering works
# The filtering is needed because other tests will put stuff int he DB
a6 = JobCalculation(**calc_params)
a6.store()
a7 = Node()
a7.store()
# Query by calculation
results = list(JobCalculation.query(dbextras__key=extra_name))
# a3, a4, a5 should not be found because they are not JobCalculations.
# a6, a7 should not be found because they have not the attribute set.
self.assertEquals(set([i.pk for i in results]), set([a1.pk, a2.pk]))
# Same query, but by the generic Node class
results = list(Node.query(dbextras__key=extra_name))
self.assertEquals(set([i.pk for i in results]),
set([a1.pk, a2.pk, a3.pk, a4.pk, a5.pk]))
# Same query, but by the Data class
results = list(Data.query(dbextras__key=extra_name))
self.assertEquals(set([i.pk for i in results]), set([a3.pk, a4.pk]))
# Same query, but by the ParameterData subclass
results = list(ParameterData.query(dbextras__key=extra_name))
self.assertEquals(set([i.pk for i in results]), set([a4.pk]))
# Same query, but by the TemplateReplacerCalc subclass
results = list(TemplateReplacerCalc.query(dbextras__key=extra_name))
self.assertEquals(set([i.pk for i in results]), set([a2.pk]))
def listfamilies(self, *args):
from aiida.backends.utils import load_dbenv, is_dbenv_loaded
import argparse
parser = argparse.ArgumentParser(
prog=self.get_full_command_name(),
description='List AiiDA PAW families.')
parser.add_argument(
'-e',
'--element',
nargs='+',
type=str,
default=[],
help="Filter the families only to those containing "
"potentials for each of the specified elements")
parser.add_argument(
'-s',
'--symbol',
nargs='+',
type=str,
default=[],
help="Filter the families only to those containing "
"a potential for each of the specified symbols")
parser.add_argument(
'-d',
'--with-description',
dest='with_description',
action='store_true',
help="Show also the description for the PAW family")
parser.set_defaults(with_description=False)
params = parser.parse_args(args)
if not is_dbenv_loaded():
load_dbenv()
from aiida.orm import DataFactory
Paw = DataFactory('vasp.paw')
groups = Paw.get_paw_groups(elements=list(params.element),
symbols=list(params.symbol))
if groups:
for g in groups:
paws = Paw.query(dbgroups=g.dbgroup).distinct()
num_paws = paws.count()
description_string = ''
if params.with_description:
description_string = ': {}'.format(g.description)
groupitem = '* {} [{} pseudos]{}'
print groupitem.format(g.name, num_paws, description_string)
else:
print 'No PAW pseudopotential family found.'
