from numbers import Number, Integral
from itertools import product
from collections import defaultdict
import copy
from scipy.special import sph_harm, erf, airy, jn, gammaln
import numpy as np
import sympy as sp
from shenfun.config import config
from shenfun.optimization import cython
from shenfun.spectralbase import BoundaryConditions
from mpi4py_fft import DistArray
__all__ = ('Expr', 'BasisFunction', 'TestFunction', 'TrialFunction', 'Function',
'Array', 'FunctionSpace')
# Define some special functions required for spherical harmonics
cot = lambda x: 1/np.tan(x)
Ynm = lambda n, m, x, y : sph_harm(m, n, y, x)
airyai = lambda x: airy(x)[0]
besselj = jn
fun = {'airyai': airyai,
'cot': cot,
'Ynm': Ynm,
'erf': erf,
'besselj': besselj,
'loggamma': gammaln
}
[docs]
def FunctionSpace(N, family='Fourier', bc=None, dtype='d', quad=None,
domain=None, scaled=None, padding_factor=1, basis=None,
dealias_direct=False, coordinates=None, **kw):
r"""Return function space
Parameters
----------
N : int
Number of quadrature points
family : str, optional
Choose one of
- ``Chebyshev`` or ``C``,
- ``Chebyshevu`` or ``U``,
- ``Legendre`` or ``L``,
- ``Fourier`` or ``F``,
- ``Laguerre`` or ``La``,
- ``Hermite`` or ``H``
- ``Jacobi`` or ``J``
- ``Ultraspherical`` or ``Q``
bc : tuple or dict, optional
See :class:`~shenfun.spectralbase.BoundaryConditions`
dtype : str or np.dtype, optional
The datatype of physical space (input to forward transforms)
quad : str, optional
Type of quadrature
* For family=Chebyshev:
- GL - Chebyshev-Gauss-Lobatto
- GC - Chebyshev-Gauss
* For family=Chebyshevu:
- GU - Chebyshevu-Gauss
- GC - Chebyshev-Gauss
* For family=Legendre:
- LG - Legendre-Gauss
- GL - Legendre-Gauss-Lobatto
* For family=Laguerre:
- LG - Laguerre-Gauss
* For family=Hermite:
- HG - Hermite-Gauss
* For family=Jacobi:
- JG - Jacobi-Gauss
* For family=Ultraspherical:
- QG - Jacobi-Gauss
domain : two-tuple of floats, optional
The computational domain
scaled : bool
Whether to use scaled basis
basis : str
Name of basis to use, if there are more than one possible basis for a given
boundary condition. For example, there are three Dirichlet bases for the
Chebyshev family: :class:`~.chebyshev.bases.Heinrichs`, :class:`~.chebyshev.bases.ShenDirichlet`
and :class:`~.chebyshev.bases.Phi1`
padding_factor : float, optional
For padding backward transform (for dealiasing)
dealias_direct : bool, optional
Use 2/3-rule dealiasing
coordinates: 2- or 3-tuple (coordinate, position vector (, sympy assumptions)), optional
Map for curvilinear coordinates. Used for curves parametrized with one
parameter in :math:`\mathbb{R}^2` or :math:`\mathbb{R}^3`.
The parameter is the first item. Second item is a tuple of length 2 or
3 (for :math:`\mathbb{R}^2` or :math:`\mathbb{R}^3`) for the Cartesian
position vector. The sympy assumptions are there to help Sympy compute
basis vectors and metrics, see :class:`~shenfun.coordinates.Coordinates`
for options.
Example circle of radius 1::
theta = sp.Symbol('x', real=True)
rv = (sp.cos(theta), sp.sin(theta))
FunctionSpace(16, 'L', coordinates=(theta, rv), domain=(0, 2*np.pi))
Examples
--------
>>> from shenfun import FunctionSpace
>>> F0 = FunctionSpace(16, 'F')
>>> C1 = FunctionSpace(32, 'C', quad='GC')
"""
par = {'padding_factor': padding_factor,
'dealias_direct': dealias_direct,
'dtype': dtype,
'coordinates': coordinates}
par.update(kw)
if domain is not None:
par['domain'] = domain
if family.lower() in ('fourier', 'f'):
from shenfun import fourier
if np.dtype(dtype).char in 'FDG':
B = fourier.bases.C2C
else:
B = fourier.bases.R2C
del par['dtype']
return B(N, **par)
elif family.lower() in ('chebyshev', 'c'):
from shenfun import chebyshev
if quad is not None:
assert quad in ('GC', 'GL', 'GU')
par['quad'] = quad
if scaled is not None:
assert isinstance(scaled, bool)
par['scaled'] = scaled
# Boundary conditions abbreviated in dictionary keys as
# left->L, right->R, Dirichlet->D, Neumann->N
# So LDRD means left Dirichlet, right Dirichlet
bases = defaultdict(lambda: chebyshev.bases.Generic,
{
'': chebyshev.bases.Orthogonal,
'LDRD': chebyshev.bases.ShenDirichlet,
'LNRN': chebyshev.bases.ShenNeumann,
'LDRN': chebyshev.bases.DirichletNeumann,
'LNRD': chebyshev.bases.NeumannDirichlet,
'RD': chebyshev.bases.UpperDirichlet,
'LD': chebyshev.bases.LowerDirichlet,
'RDRN': chebyshev.bases.UpperDirichletNeumann,
'LDLN': chebyshev.bases.LowerDirichletNeumann,
'LDLNRDRN': chebyshev.bases.ShenBiharmonic,
'LDLNLN2RDRNRN2': chebyshev.bases.Compact3,
})
if basis is not None:
assert isinstance(basis, str)
B = getattr(chebyshev.bases, basis)
if isinstance(bc, tuple) or 'Generic' in basis:
par['bc'] = bc
else:
if isinstance(bc, (str, tuple, dict)):
domain = (-1, 1) if domain is None else domain
bcs = BoundaryConditions(bc, domain=domain)
key = bcs.orderednames()
par['bc'] = bcs
elif bc is None:
key = ''
else:
raise RuntimeError
B = bases[''.join(key)]
return B(N, **par)
elif family.lower() in ('legendre', 'l'):
from shenfun import legendre
bases = defaultdict(lambda: legendre.bases.Generic,
{
'': legendre.bases.Orthogonal,
'LDRD': legendre.bases.ShenDirichlet,
'LNRN': legendre.bases.ShenNeumann,
'LDRN': legendre.bases.DirichletNeumann,
'LNRD': legendre.bases.NeumannDirichlet,
'LD': legendre.bases.LowerDirichlet,
'RD': legendre.bases.UpperDirichlet,
'RDRN': legendre.bases.UpperDirichletNeumann,
'LDLNRDRN': legendre.bases.ShenBiharmonic,
'LDLNLN2RDRNRN2': legendre.bases.Compact3,
'LDLNRN2RN3': legendre.bases.BeamFixedFree,
})
if quad is not None:
assert quad in ('LG', 'GL')
par['quad'] = quad
if scaled is not None:
assert isinstance(scaled, bool)
par['scaled'] = scaled
if basis is not None:
assert isinstance(basis, str)
B = getattr(legendre.bases, basis)
if isinstance(bc, tuple):
par['bc'] = bc
else:
if isinstance(bc, (str, tuple, dict)):
domain = (-1, 1) if domain is None else domain
bcs = BoundaryConditions(bc, domain=domain)
key = bcs.orderednames()
par['bc'] = bcs
elif bc is None:
key = ''
else:
raise RuntimeError
B = bases[''.join(key)]
return B(N, **par)
elif family.lower() in ('chebyshevu', 'u'):
from shenfun import chebyshevu
if quad is not None:
assert quad in ('GC', 'GU')
par['quad'] = quad
if scaled is not None:
assert isinstance(scaled, bool)
par['scaled'] = scaled
bases = defaultdict(lambda: chebyshevu.bases.Generic,
{
'': chebyshevu.bases.Orthogonal,
'LDRD': chebyshevu.bases.CompactDirichlet,
'LNRN': chebyshevu.bases.CompactNeumann,
'LD': chebyshevu.bases.LowerDirichlet,
'RD': chebyshevu.bases.UpperDirichlet,
'LDLNRDRN': chebyshevu.bases.CompactBiharmonic,
'LDLNLN2RDRNRN2': chebyshevu.bases.Compact3,
#'LDLNLN2N3RDRNRN2N3': chebyshevu.bases.Phi4
})
if basis is not None:
assert isinstance(basis, str)
B = getattr(chebyshevu.bases, basis)
if isinstance(bc, tuple):
par['bc'] = bc
else:
if isinstance(bc, (str, tuple, dict)):
domain = (-1, 1) if domain is None else domain
bcs = BoundaryConditions(bc, domain=domain)
key = bcs.orderednames()
par['bc'] = bcs
elif bc is None:
key = ''
else:
raise NotImplementedError
B = bases[''.join(key)]
return B(N, **par)
elif family.lower() in ('ultraspherical', 'q'):
from shenfun import ultraspherical
if quad is not None:
assert quad == 'QG'
par['quad'] = quad
if scaled is not None:
assert isinstance(scaled, bool)
par['scaled'] = scaled
bases = defaultdict(lambda: ultraspherical.bases.Generic,
{
'': ultraspherical.bases.Orthogonal,
'LDRD': ultraspherical.bases.CompactDirichlet,
'LNRN': ultraspherical.bases.CompactNeumann,
'LD': ultraspherical.bases.LowerDirichlet,
'RD': ultraspherical.bases.UpperDirichlet,
'LDLNRDRN': ultraspherical.bases.CompactBiharmonic,
#'LDLNLN2RDRNRN2': ultraspherical.bases.Phi3,
#'LDLNLN2N3RDRNRN2N3': ultraspherical.bases.Phi4
})
if basis is not None:
assert isinstance(basis, str)
B = getattr(ultraspherical.bases, basis)
if isinstance(bc, tuple):
par['bc'] = bc
else:
if isinstance(bc, (str, tuple, dict)):
domain = (-1, 1) if domain is None else domain
bcs = BoundaryConditions(bc, domain=domain)
key = bcs.orderednames()
par['bc'] = bcs
elif bc is None:
key = ''
else:
raise NotImplementedError
B = bases[''.join(key)]
return B(N, **par)
elif family.lower() in ('laguerre', 'la'):
from shenfun import laguerre
if quad is not None:
assert quad in ('LG',)
par['quad'] = quad
bases = defaultdict(lambda: laguerre.bases.Generic,
{
'': laguerre.bases.Orthogonal,
'LD': laguerre.bases.CompactDirichlet,
'LN': laguerre.bases.CompactNeumann,
})
if basis is not None:
assert isinstance(basis, str)
B = getattr(laguerre.bases, basis)
if isinstance(bc, tuple):
par['bc'] = bc
else:
if isinstance(bc, (str, tuple, dict)):
domain = (0, np.inf) if domain is None else domain
bcs = BoundaryConditions(bc, domain=domain)
key = bcs.orderednames()
par['bc'] = bcs
elif bc is None:
key = ''
else:
raise NotImplementedError
B = bases[''.join(key)]
return B(N, **par)
elif family.lower() in ('hermite', 'h'):
from shenfun import hermite
if quad is not None:
assert quad in ('HG',)
par['quad'] = quad
B = hermite.bases.Orthogonal
if isinstance(bc, tuple):
assert len(bc) == 2
par['bc'] = bc
elif isinstance(bc, str):
assert bc.lower() == 'dirichlet'
else:
assert bc is None
return B(N, **par)
elif family.lower() in ('jacobi', 'j'):
from shenfun import jacobi
if quad is not None:
assert quad in ('JG', )
par['quad'] = quad
if scaled is not None:
assert isinstance(scaled, bool)
par['scaled'] = scaled
bases = defaultdict(lambda: jacobi.bases.Generic,
{
'': jacobi.bases.Orthogonal,
'LDRD': jacobi.bases.CompactDirichlet,
'LNRN': jacobi.bases.CompactNeumann,
'UD': jacobi.bases.UpperDirichlet,
'LD': jacobi.bases.LowerDirichlet,
'LDLNRDRN': jacobi.bases.Phi2,
'LDLNLN2RDRNRN2': jacobi.bases.Phi3,
#'LDLNLN2N3RDRNRN2N3': jacobi.bases.Phi4,
})
if isinstance(bc, (str, tuple, dict)):
domain = (-1, 1) if domain is None else domain
bcs = BoundaryConditions(bc, domain=domain)
key = bcs.orderednames()
par['bc'] = bcs
elif bc is None:
key = ''
else:
raise NotImplementedError
if basis is not None:
assert isinstance(basis, str)
B = getattr(jacobi.bases, basis)
else:
B = bases[''.join(key)]
return B(N, **par)
else:
raise NotImplementedError
[docs]
class Expr:
r"""
Class for spectral Galerkin forms
An Expression that is linear in :class:`.TestFunction`,
:class:`.TrialFunction` or :class:`.Function`. Exprs are used as input
to :func:`.inner` or :func:`.project`.
Parameters
----------
basis : :class:`.BasisFunction`
:class:`.TestFunction`, :class:`.TrialFunction` or :class:`.Function`
terms : list of list of lists of length dimension
Describes the differential operations performed on the basis function
in the `Expr`
The triply nested `terms` lists are such that
- the outermost list represents a tensor component. There is one item for
each tensor component. If the Expr is a scalar with rank = 0, then
len(terms) = 1. For vectors it equals ndim, which is the number of
dimensions, and for second order tensors it equals ndim*ndim.
- the second nested list represents the different terms in the form, that
may be more than one. For example, the scalar valued `div(grad(u))` has
three terms in 3D:
.. math::
\partial^2u/\partial x^2 + \partial^2u/\partial y^2 + \partial^2u/\partial z^2
- the last inner list represents the differential operations for each term
and each tensor component, stored for each as a list of length = dim
The Expr `div(grad(u))`, where u is a scalar, is as such represented
as a nested list of shapes (1, 3, 3), 1 meaning it's a scalar, the first 3
because the Expr consists of the sum of three terms, and the last 3
because it is 3D. The entire representation is::
[[[2, 0, 0],
[0, 2, 0],
[0, 0, 2]]]
where the first [2, 0, 0] term has two derivatives in first direction
and none in the others, the second [0, 2, 0] has two derivatives in
second direction, and the last [0, 0, 2] has two derivatives in the
last direction and none in the others.
scales : list of lists
Representing a scalar multiply of each inner product. Note that
the scalar can also be a function of coordinates (using sympy).
There is one scale for each term in each tensor component, and as
such it is a list of lists.
indices : list of lists
Index into :class:`.CompositeSpace`. Only used when the basis of the
form is composite. There is one index for each term in each tensor component,
and as such it is a list of lists.
Examples
--------
>>> from shenfun import *
>>> from mpi4py import MPI
>>> comm = MPI.COMM_WORLD
>>> C0 = FunctionSpace(16, 'F', dtype='D')
>>> C1 = FunctionSpace(16, 'F', dtype='D')
>>> R0 = FunctionSpace(16, 'F', dtype='d')
>>> T = TensorProductSpace(comm, (C0, C1, R0))
>>> v = TestFunction(T)
>>> e = div(grad(v))
>>> e.terms()
[[[2, 0, 0], [0, 2, 0], [0, 0, 2]]]
>>> e2 = grad(v)
>>> e2.terms()
[[[1, 0, 0]], [[0, 1, 0]], [[0, 0, 1]]]
Note that `e2` in the example has shape (3, 1, 3). The first 3 because it
is a vector, the 1 because each vector item contains one term, and the
final 3 since it is a 3-dimensional tensor product space.
"""
def __init__(self, basis, terms=None, scales=None, indices=None):
self._basis = basis
self._terms = terms
self._scales = scales
self._indices = indices
ndim = self.function_space().dimensions
if terms is None:
self._terms = np.zeros((self.function_space().num_components(), 1, ndim),
dtype=int).tolist()
if scales is None:
self._scales = np.ones((self.function_space().num_components(), 1), dtype=object).tolist()
if indices is None:
self._indices = (basis.offset()+np.arange(self.function_space().num_components())[:, np.newaxis]).tolist()
[docs]
def basis(self):
"""Return basis of Expr"""
return self._basis
@property
def base(self):
"""Return base BasisFunction used in Expr"""
return self._basis if self._basis.base is None else self._basis.base
[docs]
def function_space(self):
"""Return function space of basis in Expr"""
return self._basis.function_space()
[docs]
def terms(self):
"""Return terms of Expr"""
return self._terms
[docs]
def scales(self):
"""Return scales of Expr"""
return self._scales
@property
def argument(self):
"""Return argument of Expr's basis"""
return self._basis.argument
[docs]
def set_basis(self, u):
"""Change the basis function this Expr applies to"""
self._basis = u
[docs]
def expr_rank(self):
"""Return rank of Expr"""
if self.dimensions == 1:
assert len(self._terms) < 3
return len(self._terms)-1
if len(self._terms) == 1:
return 0
if len(self._terms) == len(self._terms[0][0]):
return 1
if len(self._terms) == len(self._terms[0][0])**2:
return 2
raise RuntimeError
@property
def tensor_rank(self):
"""Return rank of Expr's :class:`BasisFunction`"""
return self._basis.tensor_rank
[docs]
def indices(self):
"""Return indices of Expr"""
return self._indices
[docs]
def num_components(self):
"""Return number of components in Expr"""
return len(self._terms)
[docs]
def num_terms(self):
"""Return number of terms in Expr"""
return [len(terms) for terms in self._terms]
@property
def dimensions(self):
"""Return ndim of Expr"""
return self.function_space().dimensions
#return len(self._terms[0][0])
[docs]
def index(self):
if self.num_components() == 1:
return self._basis.offset()
return None
[docs]
def tolatex(self, symbol_names=None, funcname='u', replace=None):
s = ""
x = 'xyzrst'
xx = list(product(x[:self.dimensions], x[:self.dimensions]))
symbols = {k: k for k in x[:self.dimensions]}
if symbol_names is not None:
symbols = {str(k): val for k, val in symbol_names.items()}
bl = 'e' if config['basisvectors'] == 'normal' else 'b'
for i, vec in enumerate(self.terms()):
if self.num_terms()[i] == 0 and self.num_components() > 1:
s += '0 \\mathbf{%s}_{%s} \\\\ +'%(bl, symbols[x[i]])
continue
if self.num_components() > 1:
s += '\\left( '
for j, term in enumerate(vec):
sc = self.scales()[i][j]
k = self.indices()[i][j]
kk = k % self.dimensions
if not sc == 1:
if replace is not None:
for repl in replace:
assert len(repl) == 2
sc = sc.replace(*repl)
#sc = sp.simplify(sc)
scp = sp.latex(sc, symbol_names=symbol_names)
if isinstance(sc, sp.Add):
scp = "\\left( %s \\right)"%(scp)
elif scp.startswith('-'):
s = s.rstrip('+')
if scp == '-1':
scp = '-'
s += scp
t = np.array(term)
cmp = funcname
if self.tensor_rank == 1:
cmp = funcname + '^{%s}'%(symbols[x[kk]])
elif self.tensor_rank == 2:
cmp = funcname + '^{%s%s}'%(tuple([symbols[xj] for xj in xx[k]]))
if np.sum(t) == 0:
s += cmp
else:
p = '^'+str(np.sum(t)) if np.sum(t) > 1 else ' '
s += "\\frac{\\partial%s %s}{"%(p, cmp)
for j, ti in enumerate(t):
if ti > 0:
tt = '^'+str(ti) if ti > 1 else ' '
s += "\\partial %s%s "%(symbols[x[j]], tt)
s += "}"
s += '+'
if self.num_components() > 1:
s = s.rstrip('+')
if self.expr_rank() == 1:
s += "\\right) \\mathbf{%s}_{%s} \\\\"%(bl, symbols[x[i]])
elif self.expr_rank() == 2:
s += "\\right) \\mathbf{%s}_{%s} \\mathbf{%s}_{%s} \\\\"%(bl, symbols[xx[i][0]], bl, symbols[xx[i][1]])
s += '+'
return r"""%s"""%(s.rstrip('+'))
[docs]
def tosympy(self, basis=None, psi=sp.symbols('x,y,z,r,s', real=True)):
"""Return self evaluated with a sympy basis
Parameters
----------
basis : sympy Expr or string
if sympy Expr, then use this directly
if str, then use that as name for a generic sympy Function
psi : tuple of sympy Symbols
Examples
--------
>>> from shenfun import FunctionSpace, TensorProductSpace
>>> import sympy as sp
>>> theta, r = psi = sp.symbols('x,y', real=True, positive=True)
>>> rv = (r*sp.cos(theta), r*sp.sin(theta))
>>> F0 = FunctionSpace(8, 'F', dtype='d')
>>> T0 = FunctionSpace(8, 'C')
>>> T = TensorProductSpace(comm, (F0, T0), coordinates=(psi, rv))
>>> u = TrialFunction(T)
>>> ue = r*sp.cos(theta)
>>> du = div(grad(u))
>>> du.tosympy()
Derivative(u(x, y), (y, 2)) + Derivative(u(x, y), y)/y + Derivative(u(x, y), (x, 2))/y**2
>>> du.tosympy(basis=ue, psi=psi)
0
"""
ndim = self.dimensions
if basis is None:
basis = 'u'
if isinstance(basis, str):
# Create sympy Function
st = basis
if self.tensor_rank == 1:
basis = []
for i in range(self.dimensions):
basis.append(sp.Function(st+'xyzrs'[i])(*psi[:self.dimensions]))
elif self.tensor_rank == 2:
raise NotImplementedError
else:
basis = [sp.Function(st)(*psi[:self.dimensions])]
else:
if isinstance(basis, sp.Expr):
basis = [basis]
assert len(basis) == ndim**self.tensor_rank
u = []
for i, vec in enumerate(self.terms()):
s = sp.S(0)
for j, term in enumerate(vec):
sc = self.scales()[i][j]
k = self.indices()[i][j]
b0 = basis[k]
tt = tuple([psi[n] for n, l in enumerate(term) for m in range(l)])
bi = 1
if np.sum(term) > 0:
ss = sc*bi*sp.diff(b0, *tt)
s += ss
else:
ss = sc*b0*bi
s += ss
u.append(s)
if len(u) == 1:
return u[0]
return u
[docs]
def eval(self, x, output_array=None):
"""Return expression evaluated on x
Parameters
----------
x : float or array of floats
Array must be of shape (D, N), for N points in D dimensions
"""
from shenfun import CompositeSpace
from shenfun.fourier.bases import R2C
if len(x.shape) == 1: # 1D case
x = x[None, :]
V = self.function_space()
basis = self.basis()
if output_array is None:
output_array = np.zeros(x.shape[1], dtype=V.forward.input_array.dtype)
else:
output_array[:] = 0
work = np.zeros_like(output_array)
assert V.dimensions == len(x)
for vec, (base, ind) in enumerate(zip(self.terms(), self.indices())):
for base_j, b0 in enumerate(base):
M = []
test_sp = V
if isinstance(V, CompositeSpace):
test_sp = V.flatten()[ind[base_j]]
r2c = -1
last_conj_index = -1
sl = -1
for axis, k in enumerate(b0):
xx = np.atleast_1d(test_sp[axis].map_reference_domain(np.squeeze(x[axis])))
P = test_sp[axis].evaluate_basis_derivative_all(xx, k=k)
if not test_sp[axis].domain_factor() == 1:
P *= test_sp[axis].domain_factor()**(k)
if len(x) > 1:
M.append(P[..., V.local_slice()[axis]])
if isinstance(test_sp[axis], R2C) and len(x) > 1:
r2c = axis
m = test_sp[axis].N//2+1
if test_sp[axis].N % 2 == 0:
last_conj_index = m-1
else:
last_conj_index = m
sl = V.local_slice()[axis].start
bv = basis if basis.tensor_rank == 0 else basis[ind[base_j]]
work.fill(0)
if len(x) == 1:
work = np.dot(P, bv)
elif len(x) == 2:
work = cython.evaluate.evaluate_2D(work, bv, M, r2c, last_conj_index, sl)
elif len(x) == 3:
work = cython.evaluate.evaluate_3D(work, bv, M, r2c, last_conj_index, sl)
sc = self.scales()[vec][base_j]
if not hasattr(sc, 'free_symbols'):
sc = float(sc)
else:
sym0 = tuple(sc.free_symbols)
m = []
for sym in sym0:
j = 'xyzrs'.index(str(sym))
m.append(x[j])
sc = sp.lambdify(sym0, sc)(*m)
output_array += sc*work
return output_array
@property
def is_basis_function(self):
return np.all(np.array([q for l in self.terms() for s in l for q in s]) == 0)
def __getitem__(self, i):
if i >= self.dimensions:
raise IndexError
basis = self._basis
if self.tensor_rank:
if ((self.expr_rank() == self.tensor_rank and self.is_basis_function) or
(self.expr_rank() > self.tensor_rank and self.tensor_rank > 0)):
basis = self._basis[i]
if self.expr_rank() == 1:
return Expr(basis,
[copy.deepcopy(self._terms[i])],
[copy.deepcopy(self._scales[i])],
[copy.deepcopy(self._indices[i])])
elif self.expr_rank() == 2:
ndim = self.dimensions
return Expr(basis,
copy.deepcopy(self._terms[i*ndim:(i+1)*ndim]),
copy.deepcopy(self._scales[i*ndim:(i+1)*ndim]),
copy.deepcopy(self._indices[i*ndim:(i+1)*ndim]))
else:
raise NotImplementedError
[docs]
def get_contravariant_component(self, i, j=None):
"""Return contravariant vector component
Parameters
----------
i : int or None
First vector component. In 1D set i to None
j : int, optional
Second component of tensor
Note
----
The contravariant vector components are the components to the
covariant basisvectors. If the basis vectors are normalized, i.e.,
if the setting is config['basisvectors'] == 'normal', then the
normal vector components must be scaled in order to get to the
contravariant components.
"""
if i is None:
# e.g., using grad in 1D
if config['basisvectors'] == 'normal':
assert self.dimensions == 1
return self*(1/self.function_space().coors.hi[0])
assert config['basisvectors'] == 'covariant'
return self
if j is None:
if i >= self.dimensions:
raise IndexError
if config['basisvectors'] == 'normal':
if self.dimensions == 1: # e.g., using grad in 1D
assert i == 0
return self*(1/self.function_space().coors.hi[i])
return self[i]*(1/self.function_space().coors.hi[i])
assert config['basisvectors'] == 'covariant'
return self[i]
assert isinstance(j, int)
if i >= self.dimensions or j >= self.dimensions:
raise IndexError
if config['basisvectors'] == 'normal':
hi = self.function_space().coors.hi
return self[i][j]*(1/(hi[i]*hi[j]))
assert config['basisvectors'] == 'covariant'
return self[i][j]
def __mul__(self, a):
sc = copy.deepcopy(self.scales())
coors = self.function_space().coors
for i, vec in enumerate(sc):
ai = a[i] if isinstance(a, tuple) else a
ai = sp.sympify(ai)
for j in range(len(vec)):
vec[j] = coors.refine(vec[j]*ai)
vec[j] = sp.simplify(vec[j], measure=self.function_space().coors._measure)
return Expr(self._basis, copy.deepcopy(self._terms), sc, copy.deepcopy(self._indices))
def __rmul__(self, a):
return self.__mul__(a)
def __imul__(self, a):
sc = self.scales()
coors = self.function_space().coors
for i, vec in enumerate(sc):
ai = a[i] if isinstance(a, tuple) else a
ai = sp.sympify(ai)
for j in range(len(vec)):
vec[j] = coors.refine(vec[j])
vec[j] = sp.simplify(vec[j]*ai, measure=coors._measure)
return self
def __add__(self, a):
assert isinstance(a, (Expr, BasisFunction))
if not isinstance(a, Expr):
a = Expr(a)
assert self.num_components() == a.num_components()
#assert self.function_space() == a.function_space()
assert self.argument == a.argument
if id(self._basis) == id(a._basis):
basis = self._basis
else:
assert id(self.base) == id(a.base)
basis = self.base
# Concatenate terms
terms, scales, indices = [], [], []
for i in range(self.num_components()):
terms.append(self._terms[i] + a.terms()[i])
scales.append(self._scales[i] + a.scales()[i])
indices.append(self._indices[i] + a.indices()[i])
return Expr(basis, terms, scales, indices)
def __iadd__(self, a):
assert isinstance(a, (Expr, BasisFunction))
if not isinstance(a, Expr):
a = Expr(a)
assert self.num_components() == a.num_components()
assert self.argument == a.argument
self._basis = self.base
for i in range(self.num_components()):
self._terms[i] += a.terms()[i]
self._scales[i] += a.scales()[i]
self._indices[i] += a.indices()[i]
return self
def __sub__(self, a):
assert isinstance(a, (Expr, BasisFunction))
if not isinstance(a, Expr):
a = Expr(a)
assert self.num_components() == a.num_components()
#assert self.function_space() == a.function_space()
assert self.argument == a.argument
if id(self._basis) == id(a._basis):
basis = self._basis
else:
assert id(self.base) == id(a.base)
basis = self.base
# Concatenate terms
terms, scales, indices = [], [], []
for i in range(self.num_components()):
terms.append(self._terms[i] + a.terms()[i])
scales.append(self._scales[i] + (-np.array(a.scales()[i])).tolist())
indices.append(self._indices[i] + a.indices()[i])
return Expr(basis, terms, scales, indices)
def __isub__(self, a):
assert isinstance(a, (Expr, BasisFunction))
if not isinstance(a, Expr):
a = Expr(a)
assert self.num_components() == a.num_components()
assert self.function_space() == a.function_space()
assert self.argument == a.argument
if id(self._basis) == id(a._basis):
basis = self._basis
else:
assert id(self.base) == id(a.base)
basis = self.base
self._basis = basis
for i in range(self.num_components()):
self._terms[i] += a.terms()[i]
self._scales[i] += (-np.array(a.scales()[i])).tolist()
self._indices[i] += a.indices()[i]
return self
def __neg__(self):
sc = copy.deepcopy(self.scales())
for vec in sc:
for j in range(len(vec)):
vec[j] = vec[j]*(-1)
return Expr(self.basis(), copy.deepcopy(self.terms()), sc,
copy.deepcopy(self.indices()))
[docs]
def simplify(self):
"""Join terms, that are otherwise equal, using scale"""
if np.all(np.array(self.num_terms()) == 1):
return
tms = []
inds = []
scs = []
for (terms, indices, scales) in zip(self.terms(), self.indices(), self.scales()):
tms.append([])
inds.append([])
scs.append([])
for i, (term, ind, sc) in enumerate(zip(terms, indices, scales)):
match = []
if len(tms[-1]) > 0:
match = np.where((np.array(term) == np.array(tms[-1])).prod(axis=1) == 1)[0]
matchi = np.where(ind == np.array(inds[-1]))[0]
match = np.intersect1d(match, matchi)
if len(match) > 0:
k = match[0]
assert inds[-1][k] == ind
scs[-1][k] += sc
scs[-1][k] = sp.simplify(scs[-1][k], measure=self.function_space().coors._measure)
scs[-1][k] = self.function_space().coors.refine(scs[-1][k])
if scs[-1][k] == 0: # Remove if scale is zero
scs[-1].pop(k)
tms[-1].pop(k)
inds[-1].pop(k)
continue
sc = sp.simplify(sc, measure=self.function_space().coors._measure)
sc = self.function_space().coors.refine(sc)
if not sc == 0:
tms[-1].append(term)
inds[-1].append(ind)
scs[-1].append(sc)
self._terms = tms
self._indices = inds
self._scales = scs
return
def __eq__(self, a):
if self.base != a.base:
return False
if self.num_components() != a.num_components():
return False
if not np.all(self.num_terms() == a.num_terms()):
return False
for s0, q0 in zip(self.scales(), a.scales()):
for si, qi in zip(s0, q0):
if si != qi:
return False
for s0, q0 in zip(self.indices(), a.indices()):
for si, qi in zip(s0, q0):
if si != qi:
return False
for t0, p0 in zip(self.terms(), a.terms()):
for ti, pi in zip(t0, p0):
if not np.all(ti == pi):
return False
return True
[docs]
def subs(self, a, b):
"""Replace `a` with `b` in scales
Parameters
----------
a : Sympy Symbol
b : Number
"""
scs = self.scales()
for i, sci in enumerate(scs):
for j, scj in enumerate(sci):
scs[i][j] = scj.subs(a, b)
[docs]
class BasisFunction:
"""Base class for arguments to shenfun's :class:`.Expr`
Parameters
----------
space : :class:`.TensorProductSpace`, :class:`.CompositeSpace` or
:class:`.SpectralBase`
index : int
Local component of basis in :class:`.CompositeSpace`
offset : int
The number of scalar spaces (i.e., :class:`.TensorProductSpace`)
ahead of this space
base : The base :class:`BasisFunction`
"""
def __init__(self, space, index=0, offset=0, base=None):
self._space = space
self._index = index
self._offset = offset
self._base = base
@property
def tensor_rank(self):
"""Return tensor rank of basis"""
return self.function_space().tensor_rank
[docs]
def expr_rank(self):
"""Return rank of basis"""
return self.tensor_rank
[docs]
def function_space(self):
"""Return function space of BasisFunction"""
return self._space
@property
def base(self):
"""Return base """
return self._base if self._base is not None else self
@property
def argument(self):
"""Return argument of basis"""
raise NotImplementedError
[docs]
def num_components(self):
"""Return number of components in basis"""
return self.function_space().num_components()
@property
def dimensions(self):
"""Return dimensions of function space"""
return self.function_space().dimensions
[docs]
def index(self):
"""Return index into base space"""
return self._offset + self._index
[docs]
def offset(self):
"""Return offset of this basis
The offset is the number of scalar :class:`.TensorProductSpace` es ahead
of this space in a :class:`.CompositeSpace`.
"""
return self._offset
def __getitem__(self, i):
assert self.function_space().is_composite_space
base = self.base
space = self._space[i]
offset = self._offset
for k in range(i):
offset += self._space[k].num_components()
t0 = self.__class__(space, 0, offset, base)
return t0
def __mul__(self, a):
b = Expr(self)
return b*a
def __rmul__(self, a):
return self.__mul__(a)
def __imul__(self, a):
raise RuntimeError
def __add__(self, a):
assert isinstance(a, (Expr, BasisFunction))
b = Expr(self)
return b+a
def __iadd__(self, a):
raise RuntimeError
def __sub__(self, a):
assert isinstance(a, (Expr, BasisFunction))
b = Expr(self)
return b-a
def __isub__(self, a):
raise RuntimeError
[docs]
class TestFunction(BasisFunction):
"""Test function - BasisFunction with argument = 0
Parameters
----------
space: :class:`.TensorProductSpace` or :class:`.CompositeSpace`
index: int, optional
Component of basis in :class:`.CompositeSpace`
offset : int
The number of scalar spaces (i.e., :class:`.TensorProductSpace`)
ahead of this space
base : The base :class:`.TestFunction`
"""
def __init__(self, space, index=0, offset=0, base=None):
BasisFunction.__init__(self, space, index, offset, base)
@property
def argument(self):
return 0
[docs]
class TrialFunction(BasisFunction):
"""Trial function - BasisFunction with argument = 1
Parameters
----------
space: :class:`.TensorProductSpace` or :class:`.CompositeSpace`
index: int, optional
Component of basis in :class:`.CompositeSpace`
offset : int
The number of scalar spaces (i.e., :class:`.TensorProductSpace`)
ahead of this space
base : The base :class:`.TrialFunction`
"""
def __init__(self, space, index=0, offset=0, base=None):
BasisFunction.__init__(self, space, index, offset, base)
@property
def argument(self):
return 1
class ShenfunBaseArray(DistArray):
def __new__(cls, space, val=0, buffer=None, reltol=1e-12, abstol=1e-15):
if hasattr(space, 'points_and_weights'): # 1D case
if space.N == 0:
space = space.get_adaptive(fun=buffer, reltol=reltol, abstol=abstol)
if cls.__name__ == 'Function':
dtype = space.forward.output_array.dtype
shape = space.forward.output_array.shape
else:
assert cls.__name__ == 'Array'
dtype = space.forward.input_array.dtype
shape = space.forward.input_array.shape
if not space.num_components() == 1:
shape = (space.num_components(),) + shape
# if a list of sympy expressions
if isinstance(buffer, (list, tuple)):
assert len(buffer) == len(space.flatten())
sympy_buffer = buffer
buffer = Array(space)
dtype = space.forward.input_array.dtype
for i, buf0 in enumerate(sympy_buffer):
if isinstance(buf0, Number):
buffer.v[i] = buf0
elif hasattr(buf0, 'free_symbols'):
x = buf0.free_symbols.pop()
buffer.v[i] = sp.lambdify(x, buf0)(space.mesh()).astype(dtype)
if cls.__name__ == 'Function':
buf = Function(space)
buf = buffer.forward(buf)
buffer = buf
elif hasattr(buffer, 'free_symbols'):
# Evaluate sympy function on entire mesh
syms = buffer.free_symbols
if len(syms) > 0:
x = syms.pop()
buffer = sp.lambdify(x, buffer)
buf = buffer(space.mesh()).astype(space.forward.input_array.dtype)
else:
buf = buffer
buffer = Array(space)
buffer.v[:] = buf
if cls.__name__ == 'Function':
buf = Function(space)
buf = buffer.forward(buf)
buffer = buf
elif isinstance(buffer, np.ndarray):
buffer = buffer if buffer.base is None else buffer.base
val0 = val if isinstance(val, Number) else None
obj = DistArray.__new__(cls, shape, buffer=buffer, dtype=dtype,
val=val0, rank=space.is_composite_space)
obj._space = space
obj._offset = 0
if cls.__name__ == 'Function':
obj._padded_space = {}
if buffer is None and isinstance(val, (list, tuple)):
assert len(val) == len(obj)
for i, v in enumerate(val):
obj.v[i] = v
return obj
if min(space.global_shape()) == 0:
space = space.get_adaptive(fun=buffer, reltol=reltol, abstol=abstol)
if cls.__name__ == 'Function':
forward_output = True
p0 = space.forward.output_pencil
dtype = space.forward.output_array.dtype
else:
assert cls.__name__ == 'Array'
forward_output = False
p0 = space.backward.output_pencil
dtype = space.forward.input_array.dtype
# if a list of sympy expressions
if isinstance(buffer, (list, tuple)):
assert len(buffer) == len(space.flatten())
sympy_buffer = buffer
buffer = Array(space)
adtype = space.forward.input_array.dtype
if cls.__name__ == 'Function':
buff = Function(space)
mesh = space.local_mesh(True)
for i, buf0 in enumerate(sympy_buffer):
if isinstance(buf0, Number):
buffer.v[i] = buf0
elif hasattr(buf0, 'free_symbols'):
sym0 = tuple(buf0.free_symbols)
m = []
for sym in sym0:
j = 'xyzrs'.index(str(sym))
m.append(mesh[j])
buffer.v[i] = sp.lambdify(sym0, buf0, modules=['numpy', {'airyai': airyai, 'cot': cot, 'Ynm': Ynm, 'erf': erf, 'besselj': besselj}])(*m).astype(adtype)
else:
raise NotImplementedError
if cls.__name__ == 'Function':
buff = Function(space)
buff = buffer.forward(buff)
buffer = buff
# if just one sympy expression
if hasattr(buffer, 'free_symbols'):
sym0 = tuple(buffer.free_symbols)
mesh = space.local_mesh(True)
if len(sym0) > 0:
m = []
for sym in sym0:
j = 'xyzrs'.index(str(sym))
m.append(mesh[j])
buf = sp.lambdify(sym0, buffer, modules=['numpy', fun])(*m).astype(space.forward.input_array.dtype)
else:
buf = buffer
buffer = Array(space)
buffer.v[:] = buf
if cls.__name__ == 'Function':
buf = Function(space)
buf = buffer.forward(buf)
buffer = buf
global_shape = space.global_shape(forward_output)
val0 = val if isinstance(val, Number) else None
obj = DistArray.__new__(cls, global_shape,
subcomm=p0.subcomm, val=val0, dtype=dtype,
buffer=buffer, alignment=p0.axis,
rank=space.is_composite_space)
obj._space = space
obj._offset = 0
obj._rank = space.is_composite_space
obj._base = None
if isinstance(val, (list, tuple)):
for i, v in enumerate(val):
obj.v[i] = v
if cls.__name__ == 'Function':
obj._padded_space = {}
return obj
def __array_finalize__(self, obj):
if obj is None:
return
self._space = getattr(obj, '_space', None)
self._p0 = getattr(obj, '_p0', None)
self._rank = getattr(obj, '_rank', None)
self._offset = getattr(obj, '_offset', None)
self._base = getattr(obj, '_base', None)
if hasattr(obj, '_padded_space'):
self._padded_space = obj._padded_space
def function_space(self):
"""Return function space of array ``self``"""
return self._space
def index(self):
"""Return index for scalar into mixed base space"""
if self.base is None:
return None
if self.function_space().num_components() > 1:
return None
data_self = self.ctypes.data
data_base = self.base.ctypes.data
itemsize = self.itemsize
return (data_self - data_base) // (itemsize*np.prod(self.shape))
@property
def argument(self):
"""Return argument of basis"""
return 2
@property
def global_shape(self):
"""Return global shape of ``self``"""
return self.function_space().global_shape(self.forward_output)
@property
def forward_output(self):
"""Return whether ``self`` is the result of a forward transform"""
raise NotImplementedError
def __getitem__(self, i):
if self.ndim == 1:
return np.ndarray.__getitem__(self, i)
if self._space.is_composite_space and isinstance(i, Integral):
# Return view into composite Function
space = self._space[i]
offset = 0
for j in range(i):
offset += self._space[j].num_components()
ns = space.num_components()
s = slice(offset, offset+ns) if ns > 1 else offset
v0 = np.ndarray.__getitem__(self, s)
v0._space = space
v0._offset = offset + self.offset()
v0._rank = space.is_composite_space
v0._base = self.base
return v0
return np.ndarray.__getitem__(self.v, i)
def __getattribute__(self, name):
names = ('all', 'any', 'argmax', 'argmin', 'argpartition', 'argsort', 'byteswap', 'choose', 'clip', 'compress', 'cumprod', 'cumsum', 'diagonal', 'flatten', 'getfield', 'item', 'max', 'mean', 'min', 'newbyteorder', 'nonzero', 'partition', 'prod', 'ptp', 'put', 'ravel', 'repeat', 'resize', 'round', 'searchsorted', 'setflags', 'sort', 'squeeze', 'std', 'sum', 'swapaxes', 'take', 'tobytes', 'tofile', 'tolist', 'tostring', 'trace', 'transpose', 'var', 'view', 'ndim', 'flags', 'strides', 'itemsize', 'size', 'nbytes', 'dtype', 'flat', 'ctypes')
if name in names:
a = object.__getattribute__(self, '__array__')()
return np.ndarray.__getattribute__(a, name)
return object.__getattribute__(self, name)
def dim(self):
return self.function_space().dim()
def dims(self):
return self.function_space().dims()
@property
def tensor_rank(self):
return self.function_space().tensor_rank
[docs]
class Function(ShenfunBaseArray, BasisFunction):
r"""
Spectral Galerkin function for given :class:`.TensorProductSpace` or :class:`.SpectralBase`
The Function in one dimension is
.. math::
u(x) = \sum_{k \in \mathcal{K}} \hat{u}_k \psi_k(x),
where :math:`\psi_k(x)` are the trial functions and
:math:`\{\hat{u}_k\}_{k\in\mathcal{K}}` are the expansion coefficients.
Here an index set :math:`\mathcal{K}=0, 1, \ldots, N` is used
to simplify notation.
For an M+1-dimensional TensorProductSpace with coordinates
:math:`x_0, x_1, \ldots, x_M` we get
.. math::
u(x_{0}, x_{1}, \ldots, x_{M}) = \sum_{k_0 \in \mathcal{K}_0}\sum_{k_1 \in \mathcal{K}_1} \ldots \sum_{k_M \in \mathcal{K}_M} \hat{u}_{k_0, k_1, \ldots k_M} \psi_{k_0}(x_0) \psi_{k_1}(x_1) \ldots \psi_{k_M}(x_M),
where :math:`\mathcal{K}_j` is the index set for the wavenumber mesh
along axis :math:`j`.
Note that for a Cartesian mesh in 3D it would be natural to use coordinates
:math:`(x, y, z) = (x_0, x_1, x_2)` and the expansion would be the
simpler and somewhat more intuitive
.. math::
u(x, y, z) = \sum_{l \in \mathcal{K}_0}\sum_{m \in \mathcal{K}_1} \sum_{n \in \mathcal{K}_2} \hat{u}_{l, m, n} \psi_{l}(x) \psi_{m}(y) \psi_{n}(z).
The Function's values (the Numpy array) represent the :math:`\hat{u}` array.
The trial functions :math:`\psi` may differ in the different directions.
They are chosen when creating the TensorProductSpace.
Parameters
----------
space : :class:`.TensorProductSpace` or :func:`.FunctionSpace`
val : int or float
Value used to initialize array
buffer : Numpy array, :class:`.Function` or sympy `Expr`
If array it must be of correct shape.
A sympy expression is evaluated on the quadrature mesh and
forward transformed to create the buffer array.
reltol : number, optional
Relative tolerance for adaptively finding dimension of space from
the buffer.
abstol : number, optional
Absolute tolerance for adaptively finding dimension of space from
the buffer
Note
----
For more information, see `numpy.ndarray <https://docs.scipy.org/doc/numpy/reference/generated/numpy.ndarray.html>`_
Example
-------
>>> from shenfun import comm, FunctionSpace, TensorProductSpace, Function
>>> K0 = FunctionSpace(8, 'F', dtype='D')
>>> K1 = FunctionSpace(8, 'F', dtype='d')
>>> T = TensorProductSpace(comm, [K0, K1])
>>> u = Function(T)
>>> K2 = FunctionSpace(8, 'C', bc=(0, 0))
>>> T2 = TensorProductSpace(comm, [K0, K1, K2])
>>> v = Function(T2)
You get the array the space is planned for. This is probably what you
expect:
>>> u = Function(K0)
>>> print(u.shape)
(8,)
But make the TensorProductSpace change K0 and K1 inplace:
>>> T = TensorProductSpace(comm, [K0, K1], modify_spaces_inplace=True)
>>> u = Function(K0)
>>> print(u.shape)
(8, 5)
If you want to preserve K0 as a 1D function space, then use
`modify_spaces_inplace=False`, which is the default behaviour.
Note
----
The array returned will have the same shape as the arrays
`space` is planned for. So if you want a Function on a 1D
FunctionSpace, then make sure that FunctionSpace is not planned
for a TensorProductSpace.
"""
# pylint: disable=too-few-public-methods,too-many-arguments
def __init__(self, space, val=0, buffer=None, reltol=1e-12, abstol=1e-15):
BasisFunction.__init__(self, self._space, offset=0)
[docs]
def set_boundary_dofs(self):
space = self.function_space()
if space.is_composite_space:
for i, s in enumerate(space.flatten()):
bases = s.bases if hasattr(s, 'bases') else [s]
for base in bases:
if base.has_nonhomogeneous_bcs:
base.bc.set_boundary_dofs(self.v[i], True)
else:
bases = space.bases if hasattr(space, 'bases') else [space]
for base in bases:
if base.bc:
base.bc.set_boundary_dofs(self, True)
return self
@property
def forward_output(self):
return True
@property
def base(self):
return self._base if self._base is not None else self
def __call__(self, x, output_array=None):
return self.eval(x, output_array=output_array)
[docs]
def eval(self, x, output_array=None):
"""Evaluate Function at points `x`
Parameters
----------
x : float or array of floats
output_array : array, optional
Return array, function values at points
Examples
--------
>>> import sympy as sp
>>> K0 = FunctionSpace(9, 'F', dtype='D')
>>> K1 = FunctionSpace(8, 'F', dtype='d')
>>> T = TensorProductSpace(MPI.COMM_WORLD, [K0, K1], axes=(0, 1))
>>> X = T.local_mesh()
>>> x, y = sp.symbols("x,y")
>>> ue = sp.sin(2*x) + sp.cos(3*y)
>>> ul = sp.lambdify((x, y), ue, 'numpy')
>>> ua = Array(T, buffer=ul(*X))
>>> points = np.random.random((2, 4))
>>> u = ua.forward()
>>> u0 = u.eval(points).real
>>> assert np.allclose(u0, ul(*points))
"""
return self.function_space().eval(x, self, output_array)
[docs]
def backward(self, output_array=None, kind=None, mesh=None, padding_factor=None):
"""Return Function evaluated on some quadrature mesh
Parameters
----------
output_array : array, optional
Return array, Function evaluated on mesh
kind : dict or str, optional
Can be used to overload config['transforms']['kind'], using a
dictionary with family as key and values either one of the strings
- 'fast'
- 'recursive'
- 'vandermonde'
For example kind={'chebyshev': 'vandermonde'}
Note that for one-dimensional problems one can use just the string
value and no dict
mesh : str, list or functionspace, optional
- 'quadrature' - use quadrature mesh of self
- 'uniform' - use uniform mesh
- A function space with the same mesh distribution as self
padding_factor : None or number or tuple, optional
If padding_factor is a number different from 1, then perform a
padded backward transform, using a padded work array
'self._backward_work_array'. Use a tuple for different padding in
different directions.
Note
----
If `output_array` is not provided, then a new Array is created and
returned.
If `padding_factor` is used, then the returned Array is cached and
reused for subsequent calls using the same padding_factor.
"""
if padding_factor is not None and padding_factor != 1:
space = self.get_dealiased_space(padding_factor)
if output_array is not None:
assert output_array.shape == space.shape(False)
else:
# Create a work array that can be reused
# Note that the output_array is not the same shape as
# space.backward.output_array for spaces of rank > 0.
# Also, it is probably not safe to return space.backward.output_array
if hasattr(space, '_backward_work_array'):
output_array = space._backward_work_array
else:
output_array = Array(space)
space._backward_work_array = output_array
output_array = space.backward(self, output_array, kind=kind, mesh=mesh)
return output_array
space = self.function_space()
if output_array is None:
if hasattr(mesh, 'mesh'):
output_array = Array(mesh)
else:
output_array = Array(space)
output_array = space.backward(self, output_array, kind=kind, mesh=mesh)
return output_array
[docs]
def get_dealiased_space(self, padding_factor):
"""Return a function space to be used for padded transforms
Parameters
----------
padding_factor : number or tuple
The padding factor of the backward transform
A tuple can be used for padding differently in each direction.
"""
if padding_factor is None:
return self.function_space()
if np.all(abs(np.atleast_1d(padding_factor)-1) < 1e-8):
return self.function_space()
if padding_factor in self._padded_space:
return self._padded_space[padding_factor]
space = self.function_space().get_dealiased(padding_factor)
self._padded_space[padding_factor] = space
return space
[docs]
def to_ortho(self, output_array=None):
"""Project Function to orthogonal basis"""
space = self.function_space()
if output_array is None:
output_array = Function(space.get_orthogonal())
# In case of mixed space make a loop
if space.num_components() > 1:
for x, subfunction in zip(output_array, self):
x = subfunction.to_ortho(x)
return output_array
if space.dimensions > 1:
naxes = space.get_nonperiodic_axes()
axis = naxes[0]
base = space.bases[axis]
if not base.is_orthogonal:
output_array = base.to_ortho(self, output_array)
else:
output_array[:] = self
if len(naxes) > 1:
input_array = np.zeros_like(output_array.v)
for axis in naxes[1:]:
base = space.bases[axis]
input_array[:] = output_array
if not base.is_orthogonal:
output_array = base.to_ortho(input_array, output_array)
else:
output_array = space.to_ortho(self, output_array)
return output_array
[docs]
def mask_nyquist(self, mask=None):
"""Set self to have zeros in Nyquist coefficients"""
self.function_space().mask_nyquist(self, mask=mask)
[docs]
def assign(self, u_hat):
"""Assign self to u_hat of possibly different size
Parameters
----------
u_hat : Function
Function of possibly different shape than self. Must have
the same function_space
"""
from shenfun import VectorSpace
if self.ndim == 1:
assert u_hat.__class__ == self.__class__
if self.shape[0] < u_hat.shape[0]:
self.function_space()._padding_backward(self, u_hat)
elif self.shape[0] == u_hat.shape[0]:
u_hat[:] = self
elif self.shape[0] > u_hat.shape[0]:
self.function_space()._truncation_forward(self, u_hat)
return u_hat
space = self.function_space()
newspace = u_hat.function_space()
if isinstance(space, VectorSpace):
for i, self_i in enumerate(self):
u_hat[i] = self_i.assign(u_hat[i])
return u_hat
same_bases = True
for base0, base1 in zip(space.bases, newspace.bases):
if not base0.__class__ == base1.__class__:
same_bases = False
break
assert same_bases, "Can only assign on spaces with the same underlying bases"
N = []
for newbase in newspace.bases:
N.append(newbase.N)
u_hat = self.refine(N, output_array=u_hat)
return u_hat
[docs]
def refine(self, N, output_array=None):
"""Return self with new number of quadrature points
Parameters
----------
N : number or sequence of numbers
The new number of quadrature points
Note
----
If N is smaller than for self, then a truncated array
is returned. If N is greater than before, then the
returned array is padded with zeros.
"""
from shenfun.fourier.bases import R2C
from shenfun import VectorSpace
if self.ndim == 1:
assert isinstance(N, Number)
space = self.function_space()
if output_array is None:
refined_basis = space.get_refined(N)
output_array = Function(refined_basis)
output_array = self.assign(output_array)
return output_array
if isinstance(N, Number):
N = (N,)*len(self)
elif isinstance(N, (tuple, list, np.ndarray)):
assert len(N) == len(self.function_space())
space = self.function_space()
if isinstance(space, VectorSpace):
if output_array is None:
output_array = [None]*len(self)
for i, array in enumerate(self):
output_array[i] = array.refine(N, output_array=output_array[i])
if isinstance(output_array, list):
T = output_array[0].function_space()
VT = VectorSpace(T)
output_array = np.array(output_array)
output_array = Function(VT, buffer=output_array)
return output_array
axes = [bx for ax in space.axes for bx in ax]
base = space.bases[axes[0]]
global_shape = list(self.global_shape) # Global shape in spectral space
factor = N[axes[0]]/self.function_space().bases[axes[0]].N
if isinstance(base, R2C):
global_shape[axes[0]] = int((2*global_shape[axes[0]]-2)*factor)//2+1
else:
global_shape[axes[0]] = int(global_shape[axes[0]]*factor)
c1 = DistArray(global_shape,
subcomm=self.pencil.subcomm,
dtype=self.dtype,
alignment=self.alignment)
if self.global_shape[axes[0]] <= global_shape[axes[0]]:
base._padding_backward(self, c1)
else:
base._truncation_forward(self, c1)
for ax in axes[1:]:
c0 = c1.redistribute(ax)
factor = N[ax]/self.function_space().bases[ax].N
# Get a new padded array
base = space.bases[ax]
if isinstance(base, R2C):
global_shape[ax] = int(base.N*factor)//2+1
else:
global_shape[ax] = int(global_shape[ax]*factor)
c1 = DistArray(global_shape,
subcomm=c0.pencil.subcomm,
dtype=c0.dtype,
alignment=ax)
# Copy from c0 to d0
if self.global_shape[ax] <= global_shape[ax]:
base._padding_backward(c0, c1)
else:
base._truncation_forward(c0, c1)
# Reverse transfer to get the same distribution as u_hat
for ax in reversed(axes[:-1]):
c1 = c1.redistribute(ax)
if output_array is None:
refined_space = space.get_refined(N)
output_array = Function(refined_space, buffer=c1)
else:
output_array[:] = c1
return output_array
[docs]
def copy_to_flattened(self, f=None, j=(), dims=None, sl=None):
"""Copy dofs of self to a flattened array
Parameters
----------
f : array 1D, optional
The array to copy to
j : int or tuple, optional
Index into diagonal axes (Fourier). Not used if sl is provided
dims : array 1D, optional
The dim of each variable in the composite basis
sl : sequence of slices, optional
list of slices into non-diagonal axes
Note
----
All Functions have allocated Numpy arrays assuming all dofs are being
used. However, for non-orthogonal bases with boundary conditions we use
only the first N-nb dofs, where nb is the number of boundary conditions.
This function does not copy the boundary dofs.
"""
if dims is None:
dims = self.function_space()._get_ndiag_cum_dofs()
if sl is None:
sl = self.function_space()._get_ndiag_slices(j)
if f is None:
f = np.zeros(dims[-1], dtype=self.dtype)
for i, s in enumerate(sl):
f[dims[i]:dims[i+1]] = self[tuple(s)].ravel()
return f
[docs]
def copy_from_flattened(self, f, j=(), dims=None, sl=None):
"""Copy dofs to self from a flattened array
Parameters
----------
f : array 1D
The array to copy from
j : int or tuple, optional
Index into diagonal axes (Fourier)
dims : array 1D, optional
The flat dim of each variable in the composite basis
sl : sequence of slices, optional
list of slices into non-diagonal axes
Note
----
All Functions have allocated Numpy arrays assuming all dofs are being
used. However, for non-orthogonal bases with boundary conditions we use
only the first N-nb dofs, where nb is the number of boundary conditions.
This function does not copy the boundary dofs.
"""
if dims is None:
dims = self.function_space()._get_ndiag_cum_dofs()
if sl is None:
sl = self.function_space()._get_ndiag_slices(j)
for i, s in enumerate(sl):
si = tuple(s)
self[si] = f[dims[i]:dims[i+1]].reshape(self[si].shape)
return self
[docs]
class Array(ShenfunBaseArray):
r"""
Numpy array for :class:`.TensorProductSpace`
The Array is the result of a :class:`.Function` evaluated on its quadrature
mesh.
The Function represents in 1D
.. math::
u(x) = \sum_{k \in \mathcal{K}} \hat{u}_k \psi_k(x),
where :math:`\psi_k(x)` are the trial functions and
:math:`\{\hat{u}_k\}_{k\in\mathcal{K}}` are the expansion coefficients.
Here an index set :math:`\mathcal{K}=0, 1, \ldots, N` is used to
simplify notation.
For an M+1-dimensional TensorProductSpace with coordinates
:math:`x_0, x_1, \ldots, x_M` we get
.. math::
u(x_{0}, x_{1}, \ldots, x_{M}) = \sum_{k_0 \in \mathcal{K}_0}\sum_{k_1 \in \mathcal{K}_1} \ldots \sum_{k_M \in \mathcal{K}_M} \hat{u}_{k_0, k_1, \ldots k_M} \psi_{k_0}(x_0) \psi_{k_1}(x_1) \ldots \psi_{k_M}(x_M),
where :math:`\mathcal{K}_j` is the index set for the wavenumber mesh
along axis :math:`j`.
Note that for a Cartesian mesh in 3D it would be natural to use coordinates
:math:`(x, y, z) = (x_0, x_1, x_2)` and the expansion would be the
simpler and somewhat more intuitive
.. math::
u(x, y, z) = \sum_{l \in \mathcal{K}_0}\sum_{m \in \mathcal{K}_1} \sum_{n \in \mathcal{K}_2} \hat{u}_{l, m, n} \psi_{l}(x) \psi_{m}(y) \psi_{n}(z).
The Array's values (the Numpy array) represent the left hand side,
evaluated on the Cartesian quadrature mesh. With this we mean the
:math:`u(x_i, y_j, z_k)` array, where :math:`\{x_i\}_{i=0}^{N_0}`,
:math:`\{y_j\}_{j=0}^{N_1}` and :math:`\{z_k\}_{k=0}^{N_2}` represent
the mesh along the three directions. The quadrature mesh is then
.. math::
(x_i, y_j, z_k) \quad \forall \, (i, j, k) \in [0, 1, \ldots, N_0] \times [0, 1, \ldots, N_1] \times [0, 1, \ldots, N_2]
The entire spectral Galerkin function can be obtained using the
:class:`.Function` class.
Parameters
----------
space : :class:`.TensorProductSpace` or :class:`.SpectralBase`
val : int or float
Value used to initialize array
buffer : Numpy array, :class:`.Function` or sympy `Expr`
If array it must be of correct shape.
A sympy expression is evaluated on the quadrature mesh and
the result is used as buffer.
Note
----
For more information, see `numpy.ndarray <https://docs.scipy.org/doc/numpy/reference/generated/numpy.ndarray.html>`_
Examples
--------
>>> from shenfun import comm, FunctionSpace, TensorProductSpace, Array
>>> K0 = FunctionSpace(8, 'F', dtype='D')
>>> K1 = FunctionSpace(8, 'F', dtype='d')
>>> FFT = TensorProductSpace(comm, [K0, K1])
>>> u = Array(FFT)
"""
@property
def forward_output(self):
return False
[docs]
def forward(self, output_array=None, kind=None):
"""Return Function used to evaluate Array
Parameters
----------
output_array : array, optional
Expansion coefficients
kind : str, optional
- 'fast' - use fast transform if implemented
- 'vandermonde' - Use Vandermonde matrix
"""
space = self.function_space()
if output_array is None:
output_array = Function(space)
output_array = space.forward(self, output_array, kind=kind)
return output_array
[docs]
def offset(self):
"""Return offset of this basis
The offset is the number of scalar :class:`.TensorProductSpace` es ahead
of this Arrays space in a :class:`.CompositeSpace`.
"""
return self._offset
[docs]
def get_cartesian_vector(self):
"""Return self as a Cartesian vector.
Note
----
This method is only relevant for curvilinear coordinates, where
the vector uses different basis functions from the Cartesian.
"""
assert self.tensor_rank == 1
T = self.function_space()
psi = T.coors.psi
x = T.local_mesh(True)
b = T.coors.get_basis()
bij = np.array(sp.lambdify(psi, b)(*x), dtype=object)
for bi in np.ndindex(bij.shape):
if isinstance(bij[bi], int):
bij[bi] = np.full(x[0].shape, bij[bi])
bij = np.array(bij.tolist())
sl = [slice(None)]*bij.ndim
sl[1] = None
uc = np.sum(self[tuple(sl)]*bij, axis=0)
return uc