"""
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.codegen.ast import (
Attribute, CodeBlock, FunctionCall, Node, none, String,
Token, _mk_Tuple, Variable
)
from sympy.core.basic import Basic
from sympy.core.compatibility import string_types
from sympy.core.containers import Tuple
from sympy.core.expr import Expr
from sympy.core.function import Function
from sympy.core.numbers import Float, Integer
from sympy.core.sympify import sympify
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)
