diff --git a/aiida_jutools/__init__.py b/aiida_jutools/__init__.py index 8ba8ed5..0073769 100644 --- a/aiida_jutools/__init__.py +++ b/aiida_jutools/__init__.py @@ -45,3 +45,14 @@ # # The potential problem and the solution stated above, if it becomes one, applies to other JuDFTTeam plugins as well, # should they start using aiida-jutools as common codebase (aiida-fleur, aiida-spirit, aiida-spex, ...). + +def load_ipython_extension(ipython): + """ + Register ipython extensions + """ + from aiida_jutools.ipython import NotebookAiidaExport + import masci_tools + + #Register extensions from masci-tools + masci_tools.load_ipython_extension(ipython) + ipython.register_magics(NotebookAiidaExport) \ No newline at end of file diff --git a/aiida_jutools/ipython/__init__.py b/aiida_jutools/ipython/__init__.py new file mode 100644 index 0000000..92d58ae --- /dev/null +++ b/aiida_jutools/ipython/__init__.py @@ -0,0 +1,2 @@ + +from .export import NotebookAiidaExport \ No newline at end of file diff --git a/aiida_jutools/ipython/export.py b/aiida_jutools/ipython/export.py new file mode 100644 index 0000000..bfd7ddd --- /dev/null +++ b/aiida_jutools/ipython/export.py @@ -0,0 +1,62 @@ +""" +Magics class for exporting aiida nodes from a jupyer notebook environment +""" +from IPython.core.magic import magics_class, line_magic, needs_local_scope +from aiida import orm +from collections import defaultdict +import re +import yaml + +IPYTHON_VARIABLE = re.compile('_+[0-9]*') + +@magics_class +class NotebookAiidaExport: + + @line_magic + @needs_local_scope + def export_nodes(self, line, local_ns): + + args = line.split(' ', maxsplit=1) + group_name = args[0] + + if len(args) == 2: + if args[1].lower() not in ('pk', 'uuid'): + raise ValueError('Invalid value for export format (Either pk or uuid)') + export_format = args[1].lower() + else: + export_format = 'uuid' + + #Get all AiiDA nodes in the local namespace + aiida_nodes = defaultdict(set) + for name, node in local_ns.items(): + if isinstance(node, orm.Node): + aiida_nodes[node].add(name) + + aiida_nodes_one_name = {} + #Now choose the names and prefer proper variable names over _12, etc. (ipython utilities) + export_dict = defaultdict(list) + for node, names in aiida_nodes.items(): + class_name = getattr(node, 'process_class', node.__class__).__name__ + identifier = node.pk if export_format == 'pk' else node.uuid + if len(names) == 1: + name = names.pop() + export_dict[class_name].append([name, identifier]) + aiida_nodes_one_name[name] = node + elif len({name for name in names if not re.fullmatch(IPYTHON_VARIABLE, name)}) != 0: + name = {name for name in names if not re.fullmatch(IPYTHON_VARIABLE, name)}.pop() + export_dict[class_name].append([name, identifier]) + aiida_nodes_one_name[name] = node + else: + name = names.pop() + export_dict[class_name].append([name, identifier]) + aiida_nodes_one_name[name] = node + + + group, _ = orm.Group.objects.get_or_create(f'notebook-exports/{group_name}') + for name, node in aiida_nodes_one_name.items(): + node.extras['notebook-variable'] = name + node.extras['notebook-origin'] = group_name + group.add_nodes(node) + + with open(f'{group_name}-nodes.yml', 'w') as file: + yaml.dump(dict(export_dict), file)