Source code for aiida_spirit.tools.spirit_script_builder

# -*- coding: utf-8 -*-
"""
Helper class for builder run_spirit script.
"""


[docs]class PythonScriptBuilder(): """Helper class to build python scripts with correct indentation""" indentation = '' indentation_increment = 4 * ' '
[docs] def __init__(self, indentation=''): self.indentation = indentation self.body = ''
[docs] def __enter__(self): """On entering a with clause we increase the indentation.""" self.indentation += self.indentation_increment return self
[docs] def __exit__(self, exception_type, exception_value, traceback): """On exiting a with clause we decrease the indentation.""" self.indentation = self.indentation[0:-len(self.indentation_increment)] return False # return False so that exceptions inside the with block are re-raised
[docs] def __str__(self): """String representation.""" return self.body
[docs] def __iadd__(self, contents): """Operator += appends to script body and respects indentation. Ignores leading empty lines. Automatically appends '\n'. IMPORTANT: When appending multiline strings, make sure that the first non-empty line has a correct *relative* indentation to the rest of the lines. This is because the first non-empty line will be used to figure out how much leading whitespace to remove from the rest of the lines. """ split_at_newline = contents.split('\n') # Ignore the leading empty lines i = 0 while split_at_newline[i].strip() == '': i += 1 split_at_newline = split_at_newline[i:] # Use the leading whitespace of the first non empty line l = split_at_newline[0] n_leading_whitespace = len(l) - len(l.lstrip()) for l in split_at_newline: self.body += self.indentation + l[n_leading_whitespace:] + '\n' return self
[docs] def empty_line(self): """Append an empty line to the script body""" self.body += '\n'
[docs] def block(self, header): """" Start a block with a header. e.g:: with s.block("for i in range(10):"): print(i) """ self += header return self
[docs] def write(self, file): """Write script to file""" with open(file, 'w') as f: f.write(str(self))
[docs]class SpiritScriptBuilder(PythonScriptBuilder): """Helper class to build pyton scripts for the spirit api""" _method_dict = { 'llg': 'simulation.METHOD_LLG', 'mc': 'simulation.METHOD_MC' }
[docs] def method(self, key): """Set Spirit run method (i.e. LLG, MC)""" try: return self._method_dict[key.lower()] except KeyError: raise ValueError('Invalid method!')
_solver_dict = { 'depondt': 'simulation.SOLVER_DEPONDT', 'heun': 'simulation.SOLVER_HEUN', 'sib': 'simulation.SOLVER_SIB', 'rk4': 'simulation.SOLVER_SIB', 'vp': 'simulation.SOLVER_VP', 'vp_oso': 'simulation.SOLVER_VP_OSO', 'lbfgs_oso': 'simulation.SOLVER_LBFGS_OSO', 'lbfgs_atlas': 'simulation.SOLVER_LBFGS_Atlas' }
[docs] def solver(self, key): """Set spirit solver (Depodt, VP, ...)""" try: return self._solver_dict[key.lower()] except KeyError: raise ValueError('Invalid solver!')
_module_dict = { 'configuration': 'configuration', 'simulation': 'simulation', 'geometry': 'geometry', 'state': 'state', 'io': 'io' } # Modules
[docs] def module(self, key): """Import the spirit modules""" try: return self._module_dict[key.lower()] except KeyError: raise ValueError('Invalid module!')
[docs] @staticmethod def _dict_to_arg_string(dict): # pylint: disable=redefined-builtin """Format a dict as a string of keyword arguments""" return ''.join( [', {} = {}'.format(key, val) for key, val in dict.items()])
[docs] @staticmethod def list_to_arg_string(list): # pylint: disable=redefined-builtin """Format a list as a string of positional arguments""" return ''.join([', {}'.format(l) for l in list])
[docs] def _spirit_call(self, module, function_name, *args, **kwargs): """A generic call to any of the spirit api functions.""" self += '{}.{}(p_state{}{})'.format(module, function_name, self.list_to_arg_string(args), self._dict_to_arg_string(kwargs))
[docs] def import_modules(self, *args): """Imports the modules given in `*args`. If no `*args` are given, imports all modules in the `module` dict""" if len(args) > 0: for a in args: self += 'from spirit import {}'.format(a) else: for val in self._module_dict.values(): self += 'from spirit import {}'.format(val)
[docs] def state_block(self, input_file='input_created.cfg'): """Creates the state_block and thereby creates the p_state""" return self.block( "with state.State(\"{}\") as p_state:".format(input_file))
[docs] def configuration(self, fname, *args, **kwargs): """Sets one of the configuration api functions on of p_state""" self._spirit_call(self.module('configuration'), fname, *args, **kwargs)
[docs] def start_simulation(self, method, solver, *args, **kwargs): """Start a simulation with a method and a solver.""" self._spirit_call(self.module('simulation'), 'start', self.method(method), self.solver(solver), *args, **kwargs)