Source code for sympy.codegen.fnodes

"""
AST nodes specific to Fortran.

The functions defined in this module allows the user to express functions such as ``dsign``
as a SymPy function for symbolic manipulation.
"""

from sympy.core.basic import Basic
from sympy.core.expr import Expr
from sympy.core.compatibility import string_types
from sympy.core.containers import Tuple
from sympy.core.function import Function
from sympy.core.numbers import Float, Integer
from sympy.core.sympify import sympify
from sympy.codegen.ast import (
    Attribute, CodeBlock, Declaration, FunctionCall, Node, none, String,
    Token, Type, _mk_Tuple, Variable
)
from sympy.logic import true, false
from sympy.utilities.iterables import iterable



pure = Attribute('pure')
elemental = Attribute('elemental')  # (all elemental procedures are also pure)

intent_in = Attribute('intent_in')
intent_out = Attribute('intent_out')
intent_inout = Attribute('intent_inout')

allocatable = Attribute('allocatable')

[docs]class Program(Token): """ Represents a 'program' block in Fortran Examples ======== >>> from sympy.codegen.ast import Print >>> from sympy.codegen.fnodes import Program >>> prog = Program('myprogram', [Print([42])]) >>> from sympy.printing import fcode >>> print(fcode(prog, source_format='free')) program myprogram print *, 42 end program """ __slots__ = ['name', 'body'] _construct_name = String _construct_body = staticmethod(lambda body: CodeBlock(*body))
[docs]class use_rename(Token): """ Represents a renaming in a use statement in Fortran Examples ======== >>> from sympy.codegen.fnodes import use_rename, use >>> from sympy.printing import fcode >>> ren = use_rename("thingy", "convolution2d") >>> print(fcode(ren, source_format='free')) thingy => convolution2d >>> full = use('signallib', only=['snr', ren]) >>> print(fcode(full, source_format='free')) use signallib, only: snr, thingy => convolution2d """ __slots__ = ['local', 'original'] _construct_local = String _construct_original = String
def _name(arg): if hasattr(arg, 'name'): return arg.name else: return String(arg)
[docs]class use(Token): """ Represents a use statement in Fortran Examples ======== >>> from sympy.codegen.fnodes import use >>> from sympy.printing import fcode >>> fcode(use('signallib'), source_format='free') 'use signallib' >>> fcode(use('signallib', [('metric', 'snr')]), source_format='free') 'use signallib, metric => snr' >>> fcode(use('signallib', only=['snr', 'convolution2d']), source_format='free') 'use signallib, only: snr, convolution2d' """ __slots__ = ['namespace', 'rename', 'only'] defaults = {'rename': none, 'only': none} _construct_namespace = staticmethod(_name) _construct_rename = staticmethod(lambda args: Tuple(*[arg if isinstance(arg, use_rename) else use_rename(*arg) for arg in args])) _construct_only = staticmethod(lambda args: Tuple(*[arg if isinstance(arg, use_rename) else _name(arg) for arg in args]))
[docs]class Module(Token): """ Represents a module in Fortran Examples ======== >>> from sympy.codegen.fnodes import Module >>> from sympy.printing import fcode >>> print(fcode(Module('signallib', ['implicit none'], []), source_format='free')) module signallib implicit none <BLANKLINE> contains <BLANKLINE> <BLANKLINE> end module """ __slots__ = ['name', 'declarations', 'definitions'] defaults = {'declarations': Tuple()} _construct_name = String _construct_declarations = staticmethod(lambda arg: CodeBlock(*arg)) _construct_definitions = staticmethod(lambda arg: CodeBlock(*arg))
[docs]class Subroutine(Node): """ Represents a subroutine in Fortran Examples ======== >>> from sympy import symbols >>> from sympy.codegen.ast import Print >>> from sympy.codegen.fnodes import Subroutine >>> from sympy.printing import fcode >>> x, y = symbols('x y', real=True) >>> sub = Subroutine('mysub', [x, y], [Print([x**2 + y**2, x*y])]) >>> print(fcode(sub, source_format='free', standard=2003)) subroutine mysub(x, y) real*8 :: x real*8 :: y print *, x**2 + y**2, x*y end subroutine """ __slots__ = ['name', 'parameters', 'body', 'attrs'] _construct_name = String _construct_parameters = staticmethod(lambda params: Tuple(*map(Variable.deduced, params))) @classmethod def _construct_body(cls, itr): if isinstance(itr, CodeBlock): return itr else: return CodeBlock(*itr)
[docs]class SubroutineCall(Token): """ Represents a call to a subroutine in Fortran Examples ======== >>> from sympy.codegen.fnodes import SubroutineCall >>> from sympy.printing import fcode >>> fcode(SubroutineCall('mysub', 'x y'.split())) ' call mysub(x, y)' """ __slots__ = ['name', 'subroutine_args'] _construct_name = staticmethod(_name) _construct_subroutine_args = staticmethod(_mk_Tuple)
[docs]class Do(Token): """ Represents a Do loop in in Fortran Examples ======== >>> from sympy import symbols >>> from sympy.codegen.ast import aug_assign, Print >>> from sympy.codegen.fnodes import Do >>> from sympy.printing import fcode >>> i, n = symbols('i n', integer=True) >>> r = symbols('r', real=True) >>> body = [aug_assign(r, '+', 1/i), Print([i, r])] >>> do1 = Do(body, i, 1, n) >>> print(fcode(do1, source_format='free')) do i = 1, n r = r + 1d0/i print *, i, r end do >>> do2 = Do(body, i, 1, n, 2) >>> print(fcode(do2, source_format='free')) do i = 1, n, 2 r = r + 1d0/i print *, i, r end do """ __slots__ = ['body', 'counter', 'first', 'last', 'step', 'concurrent'] defaults = {'step': Integer(1), 'concurrent': false} _construct_body = staticmethod(lambda body: CodeBlock(*body)) _construct_counter = staticmethod(sympify) _construct_first = staticmethod(sympify) _construct_last = staticmethod(sympify) _construct_step = staticmethod(sympify) _construct_concurrent = staticmethod(lambda arg: true if arg else false)
[docs]class ArrayConstructor(Token): """ Represents an array constructor Examples ======== >>> from sympy.printing import fcode >>> from sympy.codegen.fnodes import ArrayConstructor >>> ac = ArrayConstructor([1, 2, 3]) >>> fcode(ac, standard=95, source_format='free') '(/1, 2, 3/)' >>> fcode(ac, standard=2003, source_format='free') '[1, 2, 3]' """ __slots__ = ['elements'] _construct_elements = staticmethod(_mk_Tuple)
[docs]class ImpliedDoLoop(Token): """ Represents an implied do loop in Fortran Examples ======== >>> from sympy import Symbol, fcode >>> from sympy.codegen.fnodes import ImpliedDoLoop, ArrayConstructor >>> i = Symbol('i', integer=True) >>> idl = ImpliedDoLoop(i**3, i, -3, 3, 2) # -27, -1, 1, 27 >>> ac = ArrayConstructor([-28, idl, 28]) # -28, -27, -1, 1, 27, 28 >>> fcode(ac, standard=2003, source_format='free') '[-28, (i**3, i = -3, 3, 2), 28]' """ __slots__ = ['expr', 'counter', 'first', 'last', 'step'] defaults = {'step': Integer(1)} _construct_expr = staticmethod(sympify) _construct_counter = staticmethod(sympify) _construct_first = staticmethod(sympify) _construct_last = staticmethod(sympify) _construct_step = staticmethod(sympify)
[docs]class Extent(Basic): """ Represents a dimension extent. Examples ======== >>> from sympy.codegen.fnodes import Extent >>> e = Extent(-3, 3) # -3, -2, -1, 0, 1, 2, 3 >>> from sympy.printing import fcode >>> fcode(e, source_format='free') '-3:3' >>> from sympy.codegen.ast import Variable, real >>> from sympy.codegen.fnodes import dimension, intent_out >>> dim = dimension(e, e) >>> arr = Variable('x', real, attrs=[dim, intent_out]) >>> fcode(arr.as_Declaration(), source_format='free', standard=2003) 'real*8, dimension(-3:3, -3:3), intent(out) :: x' """ def __new__(cls, *args): if len(args) == 2: low, high = args return Basic.__new__(cls, sympify(low), sympify(high)) elif len(args) == 0 or (len(args) == 1 and args[0] in (':', None)): return Basic.__new__(cls) # assumed shape else: raise ValueError("Expected 0 or 2 args (or one argument == None or ':')") def _sympystr(self, printer): if len(self.args) == 0: return ':' return '%d:%d' % self.args
assumed_extent = Extent() # or Extent(':'), Extent(None)
[docs]def dimension(*args): """ Creates a 'dimension' Attribute with (up to 7) extents. Examples ======== >>> from sympy.printing import fcode >>> from sympy.codegen.fnodes import dimension, intent_in >>> dim = dimension('2', ':') # 2 rows, runtime determined number of columns >>> from sympy.codegen.ast import Variable, integer >>> arr = Variable('a', integer, attrs=[dim, intent_in]) >>> fcode(arr.as_Declaration(), source_format='free', standard=2003) 'integer*4, dimension(2, :), intent(in) :: a' """ if len(args) > 7: raise ValueError("Fortran only supports up to 7 dimensional arrays") parameters = [] for arg in args: if isinstance(arg, Extent): parameters.append(arg) elif isinstance(arg, string_types): if arg == ':': parameters.append(Extent()) else: parameters.append(String(arg)) elif iterable(arg): parameters.append(Extent(*arg)) else: parameters.append(sympify(arg)) if len(args) == 0: raise ValueError("Need at least one dimension") return Attribute('dimension', parameters)
assumed_size = dimension('*')
[docs]def array(symbol, dim, intent=None, **kwargs): """ Convenience function for creating a Variable instance for a Fortran array Parameters ========== symbol : symbol dim : Attribute or iterable If dim is an ``Attribute`` it need to have the name 'dimension'. If it is not an ``Attribute``, then it is passsed to :func:`dimension` as ``*dim`` intent : str One of: 'in', 'out', 'inout' or None \\*\\*kwargs: Keyword arguments for ``Variable`` ('type' & 'value') Examples ======== >>> from sympy.printing import fcode >>> from sympy.codegen.ast import integer, real >>> from sympy.codegen.fnodes import array >>> arr = array('a', '*', 'in', type=integer) >>> print(fcode(arr.as_Declaration(), source_format='free', standard=2003)) integer*4, dimension(*), intent(in) :: a >>> x = array('x', [3, ':', ':'], intent='out', type=real) >>> print(fcode(x.as_Declaration(value=1), source_format='free', standard=2003)) real*8, dimension(3, :, :), intent(out) :: x = 1 """ if isinstance(dim, Attribute): if str(dim.name) != 'dimension': raise ValueError("Got an unexpected Attribute argument as dim: %s" % str(dim)) else: dim = dimension(*dim) attrs=list(kwargs.pop('attrs', [])) + [dim] if intent is not None: if intent not in (intent_in, intent_out, intent_inout): intent = {'in': intent_in, 'out': intent_out, 'inout': intent_inout}[intent] attrs.append(intent) value = kwargs.pop('value', None) type_ = kwargs.pop('type', None) if type_ is None: return Variable.deduced(symbol, value=value, attrs=attrs) else: return Variable(symbol, type_, value=value, attrs=attrs)
def _printable(arg): return String(arg) if isinstance(arg, string_types) else sympify(arg)
[docs]def allocated(array): """ Creates an AST node for a function call to Fortran's "allocated(...)" Examples ======== >>> from sympy.printing import fcode >>> from sympy.codegen.fnodes import allocated >>> alloc = allocated('x') >>> fcode(alloc, source_format='free') 'allocated(x)' """ return FunctionCall('allocated', [_printable(array)])
[docs]def lbound(array, dim=None, kind=None): """ Creates an AST node for a function call to Fortran's "lbound(...)" Parameters ========== array : Symbol or String dim : expr kind : expr Examples ======== >>> from sympy.printing import fcode >>> from sympy.codegen.fnodes import lbound >>> lb = lbound('arr', dim=2) >>> fcode(lb, source_format='free') 'lbound(arr, 2)' """ return FunctionCall( 'lbound', [_printable(array)] + ([_printable(dim)] if dim else []) + ([_printable(kind)] if kind else []) )
def ubound(array, dim=None, kind=None): return FunctionCall( 'ubound', [_printable(array)] + ([_printable(dim)] if dim else []) + ([_printable(kind)] if kind else []) )
[docs]def shape(source, kind=None): """ Creates an AST node for a function call to Fortran's "shape(...)" Parameters ========== source : Symbol or String kind : expr Examples ======== >>> from sympy.printing import fcode >>> from sympy.codegen.fnodes import shape >>> shp = shape('x') >>> fcode(shp, source_format='free') 'shape(x)' """ return FunctionCall( 'shape', [_printable(source)] + ([_printable(kind)] if kind else []) )
[docs]def size(array, dim=None, kind=None): """ Creates an AST node for a function call to Fortran's "size(...)" Examples ======== >>> from sympy import Symbol >>> from sympy.printing import fcode >>> from sympy.codegen.ast import FunctionDefinition, real, Return, Variable >>> from sympy.codegen.fnodes import array, sum_, size >>> a = Symbol('a', real=True) >>> body = [Return((sum_(a**2)/size(a))**.5)] >>> arr = array(a, dim=[':'], intent='in') >>> fd = FunctionDefinition(real, 'rms', [arr], body) >>> print(fcode(fd, source_format='free', standard=2003)) real*8 function rms(a) real*8, dimension(:), intent(in) :: a rms = sqrt(sum(a**2)*1d0/size(a)) end function """ return FunctionCall( 'size', [_printable(array)] + ([_printable(dim)] if dim else []) + ([_printable(kind)] if kind else []) )
[docs]def reshape(source, shape, pad=None, order=None): """ Creates an AST node for a function call to Fortran's "reshape(...)" Parameters ========== source : Symbol or String shape : ArrayExpr """ return FunctionCall( 'reshape', [_printable(source), _printable(shape)] + ([_printable(pad)] if pad else []) + ([_printable(order)] if pad else []) )
[docs]def bind_C(name=None): """ Creates an Attribute ``bind_C`` with a name Parameters ========== name : str Examples ======== >>> from sympy import Symbol >>> from sympy.printing import fcode >>> from sympy.codegen.ast import FunctionDefinition, real, Return, Variable >>> from sympy.codegen.fnodes import array, sum_, size, bind_C >>> a = Symbol('a', real=True) >>> s = Symbol('s', integer=True) >>> arr = array(a, dim=[s], intent='in') >>> body = [Return((sum_(a**2)/s)**.5)] >>> fd = FunctionDefinition(real, 'rms', [arr, s], body, attrs=[bind_C('rms')]) >>> print(fcode(fd, source_format='free', standard=2003)) real*8 function rms(a, s) bind(C, name="rms") real*8, dimension(s), intent(in) :: a integer*4 :: s rms = sqrt(sum(a**2)/s) end function """ return Attribute('bind_C', [String(name)] if name else [])
[docs]class GoTo(Token): """ Represents a goto statement in Fortran Examples ======== >>> from sympy.codegen.fnodes import GoTo >>> go = GoTo([10, 20, 30], 'i') >>> from sympy.printing import fcode >>> fcode(go, source_format='free') 'go to (10, 20, 30), i' """ __slots__ = ['labels', 'expr'] defaults = {'expr': none} _construct_labels = staticmethod(_mk_Tuple) _construct_expr = staticmethod(sympify)
[docs]class FortranReturn(Token): """ AST node explicitly mapped to a fortran "return". Because a return statement in fortran is different from C, and in order to aid reuse of our codegen ASTs the ordinary ``.codegen.ast.Return`` is interpreted as assignment to the result variable of the function. If one for some reason needs to generate a fortran RETURN statement, this node should be used. Examples ======== >>> from sympy.codegen.fnodes import FortranReturn >>> from sympy.printing import fcode >>> fcode(FortranReturn('x')) ' return x' """ __slots__ = ['return_value'] defaults = {'return_value': none} _construct_return_value = staticmethod(sympify)
class FFunction(Function): _required_standard = 77 def _fcode(self, printer): name = self.__class__.__name__ if printer._settings['standard'] < self._required_standard: raise NotImplementedError("%s requires Fortran %d or newer" % (name, self._required_standard)) return '{0}({1})'.format(name, ', '.join(map(printer._print, self.args))) class F95Function(FFunction): _required_standard = 95
[docs]class isign(FFunction): """ Fortran sign intrinsic for integer arguments. """ nargs = 2
[docs]class dsign(FFunction): """ Fortran sign intrinsic for double precision arguments. """ nargs = 2
[docs]class cmplx(FFunction): """ Fortran complex conversion function. """ nargs = 2 # may be extended to (2, 3) at a later point
[docs]class kind(FFunction): """ Fortran kind function. """ nargs = 1
[docs]class merge(F95Function): """ Fortran merge function """ nargs = 3
class _literal(Float): _token = None _decimals = None def _fcode(self, printer, *args, **kwargs): mantissa, sgnd_ex = ('%.{0}e'.format(self._decimals) % self).split('e') mantissa = mantissa.strip('0').rstrip('.') ex_sgn, ex_num = sgnd_ex[0], sgnd_ex[1:].lstrip('0') ex_sgn = '' if ex_sgn == '+' else ex_sgn return (mantissa or '0') + self._token + ex_sgn + (ex_num or '0')
[docs]class literal_sp(_literal): """ Fortran single precision real literal """ _token = 'e' _decimals = 9
[docs]class literal_dp(_literal): """ Fortran double precision real literal """ _token = 'd' _decimals = 17
class sum_(Token, Expr): __slots__ = ['array', 'dim', 'mask'] defaults = {'dim': none, 'mask': none} _construct_array = staticmethod(sympify) _construct_dim = staticmethod(sympify) class product_(Token, Expr): __slots__ = ['array', 'dim', 'mask'] defaults = {'dim': none, 'mask': none} _construct_array = staticmethod(sympify) _construct_dim = staticmethod(sympify)