# -*- coding: utf-8 -*-
###############################################################################
# Copyright (c), Forschungszentrum Jülich GmbH, IAS-1/PGI-1, Germany. #
# All rights reserved. #
# This file is part of the AiiDA-FLEUR package. #
# #
# The code is hosted on GitHub at https://github.com/JuDFTteam/aiida-fleur #
# For further information on the license, see the LICENSE.txt file #
# For further information please visit http://www.flapw.de or #
# http://aiida-fleur.readthedocs.io/en/develop/ #
###############################################################################
"""
In here we put all things util (methods, code snipets) that are often useful, but not yet in AiiDA
itself.
So far it contains:
export_extras
import_extras
delete_nodes (FIXME)
delete_trash (FIXME)
create_group
"""
# TODO import, export of descriptions, and labels...?
from __future__ import absolute_import
from __future__ import print_function
import json
import six
from six.moves import input as input_six
from aiida.orm import load_node
from aiida.orm.querybuilder import QueryBuilder
from aiida.orm import Group, Node
from aiida.common.exceptions import NotExistent
#print(extras)
'''
# uncommented this, since it is for sure out of date and dangerous to use
# for current version see aiida-core docs
def delete_nodes(pks_to_delete):
"""
Delete a set of nodes. (From AiiDA cockbook)
Note: TODO this has to be improved for workchain removal. (checkpoints and co)
Also you will be backchecked.
BE VERY CAREFUL!
.. note::
The script will also delete
all children calculations generated from the specified nodes.
:params pks_to_delete: a list of the PKs of the nodes to delete
"""
# TODO: CHECK IF THIS IS UP TO DATE!
from django.db import transaction
from django.db.models import Q
from aiida.backends.djsite.db import models
# Delete also all children of the given calculations
# Here I get a set of all pks to actually delete, including
# all children nodes.
all_pks_to_delete = set(pks_to_delete)
for pk in pks_to_delete:
all_pks_to_delete.update(
models.DbNode.objects.filter(input_links__in=pks_to_delete).values_list('pk', flat=True))
print(('I am going to delete {} nodes, including ALL THE CHILDREN'
'of the nodes you specified. Do you want to continue? [y/N]'
''.format(len(all_pks_to_delete))))
answer = input_six()
if answer.strip().lower() == 'y':
# Recover the list of folders to delete before actually deleting
# the nodes. I will delete the folders only later, so that if
# there is a problem during the deletion of the nodes in
# the DB, I don't delete the folders
# There seem to be no folders in AiiDA 1.0
#folders = [load_node(pk).folder for pk in all_pks_to_delete]
with transaction.atomic():
# Delete all links pointing to or from a given node
models.DbLink.objects.filter(Q(input__in=all_pks_to_delete) | Q(output__in=all_pks_to_delete)).delete()
# now delete nodes
models.DbNode.objects.filter(pk__in=all_pks_to_delete).delete()
# If we are here, we managed to delete the entries from the DB.
# I can now delete the folders
# There seem to be no folders in AiiDA 1.0
# for f in folders:
# f.erase()
def delete_trash():
"""
This method deletes all AiiDA nodes in the DB, which have a extra trash=True
And all their children. Could be advanced to a garbage collector.
Be careful to use it.
"""
#query db for marked trash
q = QueryBuilder()
nodes_to_delete_pks = []
q.append(Node, filters={'extras.trash': {'==': True}})
res = q.all()
for node in res:
nodes_to_delete_pks.append(node[0].pk)
print(('pk {}, extras {}'.format(node[0].pk, node[0].extras)))
#Delete the trash nodes
print(('deleting nodes {}'.format(nodes_to_delete_pks)))
delete_nodes(nodes_to_delete_pks)
return
'''
[docs]def create_group(name, nodes, description=None, add_if_exist=False):
"""
Creates a group for a given node list.
So far this is only an AiiDA verdi command.
:params name: string name for the group
:params nodes: list of AiiDA nodes, pks, or uuids
:params description: optional string that will be stored as description for the group
:returns: the group, AiiDa group
Usage example:
.. code-block:: python
group_name = 'delta_structures_gustav'
nodes_to_group_pks =[2142, 2084]
create_group(group_name, nodes_to_group_pks,
description='delta structures added by hand. from Gustavs inpgen files')
"""
#from aiida.common import NotExistent
group, created = Group.objects.get_or_create(label=name)
if created:
print(('Group created with PK={} and name {}'.format(group.pk, group.label)))
else:
print(('Group with name {} and pk {} already exists.' ''.format(group.label, group.pk)))
if add_if_exist:
print('Adding nodes to the existing group {}'.format(group.label))
else:
print('Nodes were not added to the existing group {}'.format(group.label))
return
nodes2 = []
for node in nodes:
if not isinstance(node, Node):
try:
node = load_node(node)
except NotExistent:
print('Skipping {}, it does not exist in the DB'.format(node))
continue
nodes2.append(node)
group.add_nodes(nodes2)
print(('added nodes: {} to group {} {}'.format([x.pk for x in nodes2], group.label, group.pk)))
if description:
group.description = description
return group
[docs]def get_nodes_from_group(group, return_format='uuid'):
"""
Returns a list of pk or uuid of a nodes in a given group. Since 1.1.0, this function does
not load a group using the label or any other identification. Use Group.objects.get(filter=ID) to
pre-load this, available filters are: id, uuid, label, type_string, time, description, user_id.
"""
if return_format == 'uuid':
return [x.uuid for x in group.nodes]
if return_format == 'pk':
return [x.pk for x in group.nodes]
else:
raise ValueError("return_format should be 'uuid' or 'pk'.")