diff --git a/chapman_enskog/__init__.py b/chapman_enskog/__init__.py
index ea6f81742bb7f360c880da4d3e25a7d9bb85ebe8..abe7e69bdfe7f1f3fb11ef5e331d292b2548764b 100644
--- a/chapman_enskog/__init__.py
+++ b/chapman_enskog/__init__.py
@@ -1,4 +1,4 @@
from lbmpy.chapman_enskog.derivative import DiffOperator, Diff, expandUsingLinearity, expandUsingProductRule, \
normalizeDiffOrder, chapmanEnskogDerivativeExpansion, chapmanEnskogDerivativeRecombination
-from lbmpy.chapman_enskog.chapman_enskog import getExpandedName, LbMethodEqMoments, insertMoments, takeMoments, \
+from lbmpy.chapman_enskog.chapman_enskog import LbMethodEqMoments, insertMoments, takeMoments, \
CeMoment, chainSolveAndSubstitute, timeDiffSelector, momentSelector, ChapmanEnskogAnalysis
diff --git a/chapman_enskog/chapman_enskog.py b/chapman_enskog/chapman_enskog.py
index 6883c4efa22a4ce7227c20621ff6a8d0eb8de92f..286cbfff166fc675c943c341c2727cc4a18b64e1 100644
--- a/chapman_enskog/chapman_enskog.py
+++ b/chapman_enskog/chapman_enskog.py
@@ -8,7 +8,8 @@ from lbmpy.chapman_enskog import Diff, expandUsingLinearity, expandUsingProductR
from lbmpy.chapman_enskog import DiffOperator, normalizeDiffOrder, chapmanEnskogDerivativeExpansion, \
chapmanEnskogDerivativeRecombination
from lbmpy.chapman_enskog.derivative import collectDerivatives, createNestedDiff
-from lbmpy.moments import discreteMoment, momentMatrix, polynomialToExponentRepresentation, getMomentIndices
+from lbmpy.moments import discreteMoment, momentMatrix, polynomialToExponentRepresentation, getMomentIndices, \
+ nonAliasedMoment
from pystencils.cache import diskcache
from pystencils.sympyextensions import normalizeProduct, multidimensionalSummation, kroneckerDelta
from pystencils.sympyextensions import productSymmetric
@@ -17,14 +18,6 @@ from pystencils.sympyextensions import productSymmetric
# --------------------------------------------- Helper Functions -------------------------------------------------------
-def getExpandedName(originalObject, number):
- import warnings
- warnings.warn("Deprecated!")
- name = originalObject.name
- newName = name + "^{(%i)}" % (number,)
- return originalObject.func(newName)
-
-
def expandedSymbol(name, subscript=None, superscript=None, **kwargs):
if subscript is not None:
name += "_{%s}" % (subscript,)
@@ -32,49 +25,10 @@ def expandedSymbol(name, subscript=None, superscript=None, **kwargs):
name += "^{(%s)}" % (superscript,)
return sp.Symbol(name, **kwargs)
-# -------------------------------- Summation Convention -------------------------------------------------------------
-
-
-def getOccurrenceCountOfIndex(term, index):
- if isinstance(term, Diff):
- return getOccurrenceCountOfIndex(term.arg, index) + (1 if term.target == index else 0)
- elif isinstance(term, sp.Symbol):
- return 1 if term.name.endswith("_" + str(index)) else 0
- else:
- return 0
-
-
-def replaceIndex(term, oldIndex, newIndex):
- if isinstance(term, Diff):
- newArg = replaceIndex(term.arg, oldIndex, newIndex)
- newLabel = newIndex if term.target == oldIndex else term.target
- return Diff(newArg, newLabel, term.superscript)
- elif isinstance(term, sp.Symbol):
- if term.name.endswith("_" + str(oldIndex)):
- baseName = term.name[:-(len(str(oldIndex))+1)]
- return sp.Symbol(baseName + "_" + str(newIndex))
- else:
- return term
- else:
- newArgs = [replaceIndex(a, oldIndex, newIndex) for a in term.args]
- return term.func(*newArgs) if newArgs else term
-
-# Problem: when there are more than two repeated indices... which one to replace?
# ----------------------------------------------------------------------------------------------------------------------
-def momentAliasing(momentTuple):
- moment = list(momentTuple)
- result = []
- for element in moment:
- if element > 2:
- result.append(2 - (element % 2))
- else:
- result.append(element)
- return tuple(result)
-
-
class CeMoment(sp.Symbol):
def __new__(cls, name, *args, **kwds):
obj = CeMoment.__xnew_cached_(cls, name, *args, **kwds)
@@ -256,7 +210,7 @@ def takeMoments(eqn, pdfToMomentName=(('f', '\Pi'), ('\Omega f', '\\Upsilon')),
momentTuple = tuple(momentTuple)
if useOneNeighborhoodAliasing:
- momentTuple = momentAliasing(momentTuple)
+ momentTuple = nonAliasedMoment(momentTuple)
result = CeMoment(fIndex.momentName, momentTuple, fIndex.superscript)
if derivativeTerm is not None:
result = derivativeTerm.changeArgRecursive(result)
diff --git a/moments.py b/moments.py
index ce8d32f2386b5addb9b476e91bfcafce2f9816ae..a49954089f4819c07dcb9dac4138e42a960586d3 100644
--- a/moments.py
+++ b/moments.py
@@ -1,5 +1,5 @@
+# -*- coding: utf-8 -*-
"""
-
Module Overview
~~~~~~~~~~~~~~~
@@ -25,9 +25,9 @@ All moment polynomials have to use ``MOMENT_SYMBOLS`` (which is a module variabl
Example ::
- from lbmpy.moments import MOMENT_SYMBOLS
- x, y, z = MOMENT_SYMBOLS
- secondOrderMoment = x*y + y*z
+ >>> from lbmpy.moments import MOMENT_SYMBOLS
+ >>> x, y, z = MOMENT_SYMBOLS
+ >>> secondOrderMoment = x*y + y*z
Functions
@@ -36,65 +36,17 @@ Functions
"""
import itertools
import math
-from collections import Counter, defaultdict
from copy import copy
-
+from collections import Counter, defaultdict
+from typing import Iterable, List, Sequence, Tuple, TypeVar, Optional
import sympy as sp
from pystencils.cache import memorycache
from pystencils.sympyextensions import removeHigherOrderTerms
-MOMENT_SYMBOLS = sp.symbols("x y z")
-
-def __uniqueList(seq):
- seen = {}
- result = []
- for item in seq:
- if item in seen:
- continue
- seen[item] = 1
- result.append(item)
- return result
-
-
-def __uniquePermutations(elements):
- if len(elements) == 1:
- yield (elements[0],)
- else:
- unique_elements = set(elements)
- for first_element in unique_elements:
- remaining_elements = list(elements)
- remaining_elements.remove(first_element)
- for sub_permutation in __uniquePermutations(remaining_elements):
- yield (first_element,) + sub_permutation
-
-
-def __generateFixedSumTuples(tupleLength, tupleSum, allowedValues=None, ordered=False):
- if not allowedValues:
- allowedValues = list(range(0, tupleSum + 1))
-
- assert (0 in allowedValues)
-
- def recursive_helper(currentList, position, totalSum):
- newPosition = position + 1
- if newPosition < len(currentList):
- for i in allowedValues:
- currentList[position] = i
- newSum = totalSum - i
- if newSum < 0:
- continue
- for item in recursive_helper(currentList, newPosition, newSum):
- yield item
- else:
- if totalSum in allowedValues:
- currentList[-1] = totalSum
- if not ordered:
- yield tuple(currentList)
- if ordered and currentList == sorted(currentList, reverse=True):
- yield tuple(currentList)
-
- return recursive_helper([0] * tupleLength, 0, tupleSum)
+MOMENT_SYMBOLS = sp.symbols('x y z')
+T = TypeVar('T')
# ------------------------------ Discrete (Exponent Tuples) ------------------------------------------------------------
@@ -128,12 +80,12 @@ def pickRepresentativeMoments(moments):
def momentPermutations(exponentTuple):
"""Returns all (unique) permutations of the given tuple"""
- return __uniquePermutations(exponentTuple)
+ return __unique_permutations(exponentTuple)
def momentsOfOrder(order, dim=3, includePermutations=True):
"""All tuples of length 'dim' which sum equals 'order'"""
- for item in __generateFixedSumTuples(dim, order, ordered=not includePermutations):
+ for item in __fixed_sum_tuples(dim, order, ordered=not includePermutations):
assert(len(item) == dim)
assert(sum(item) == order)
yield item
@@ -155,7 +107,7 @@ def extendMomentsWithPermutations(exponentTuples):
allMoments = []
for i in exponentTuples:
allMoments += list(momentPermutations(i))
- return __uniqueList(allMoments)
+ return __unique(allMoments)
# ------------------------------ Representation Conversions ------------------------------------------------------------
@@ -278,8 +230,7 @@ def getExponentTupleFromIndices(momentIndices, dim):
def getOrder(moment):
- """
- Computes polynomial order of given moment
+ """Computes polynomial order of given moment.
Examples:
>>> x , y, z = MOMENT_SYMBOLS
@@ -294,9 +245,32 @@ def getOrder(moment):
return sum(moment)
if len(moment.atoms(sp.Symbol)) == 0:
return 0
- leadingCoefficient = sp.polys.polytools.LM(moment)
- symbolsInLeadingCoefficient = leadingCoefficient.atoms(sp.Symbol)
- return sum([sp.degree(leadingCoefficient, gen=m) for m in symbolsInLeadingCoefficient])
+ leading_coefficient = sp.polys.polytools.LM(moment)
+ symbols_in_leading_coefficient = leading_coefficient.atoms(sp.Symbol)
+ return sum([sp.degree(leading_coefficient, gen=m) for m in symbols_in_leading_coefficient])
+
+
+def nonAliasedMoment(moment_tuple: Sequence[int]) -> Tuple[int, ...]:
+ """Takes a moment exponent tuple and returns the non-aliased version of it.
+
+ For first neighborhood stencils, all moments with exponents 3, 5, 7... are equal to exponent 1,
+ and exponents 4, 6, 8... are equal to exponent 2. This is because first neighborhood stencils only have values
+ d ∈ {+1, 0, -1}. So for example d**5 is always the same as d**3 and d, and d**6 == d**4 == d**2
+
+ Example:
+ >>> nonAliasedMoment((5, 4, 2))
+ (1, 2, 2)
+ >>> nonAliasedMoment((9, 1, 2))
+ (1, 1, 2)
+ """
+ moment = list(moment_tuple)
+ result = []
+ for element in moment:
+ if element > 2:
+ result.append(2 - (element % 2))
+ else:
+ result.append(element)
+ return tuple(result)
def isShearMoment(moment):
@@ -605,3 +579,93 @@ def monomialToPolynomialTransformationMatrix(monomials, polynomials):
exponentTuple = exponentTuple[:dim]
result[polynomialIdx, monomials.index(exponentTuple)] = factor
return result
+
+
+# --------------------------------------- Internal Functions -----------------------------------------------------------
+
+def __unique(seq: Sequence[T]) -> List[T]:
+ """Removes duplicates from a sequence in an order preserving way.
+
+ Example:
+ >>> __unique((1, 2, 3, 1, 4, 6, 3))
+ [1, 2, 3, 4, 6]
+ """
+ seen = {}
+ result = []
+ for item in seq:
+ if item in seen:
+ continue
+ seen[item] = 1
+ result.append(item)
+ return result
+
+
+def __unique_permutations(elements: Sequence[T]) -> Iterable[T]:
+ """Generates all unique permutations of the passed sequence.
+
+ Example:
+ >>> list(__unique_permutations([1, 1, 2]))
+ [(1, 1, 2), (1, 2, 1), (2, 1, 1)]
+
+ """
+ if len(elements) == 1:
+ yield (elements[0],)
+ else:
+ unique_elements = set(elements)
+ for first_element in unique_elements:
+ remaining_elements = list(elements)
+ remaining_elements.remove(first_element)
+ for sub_permutation in __unique_permutations(remaining_elements):
+ yield (first_element,) + sub_permutation
+
+
+def __fixed_sum_tuples(tuple_length: int, tuple_sum: int,
+ allowed_values: Optional[Sequence[int]] = None,
+ ordered: bool = False) -> Iterable[Tuple[int, ...]]:
+ """Generates all possible tuples of positive integers with a fixed sum of all entries.
+
+ Args:
+ tuple_length: length of the returned tuples
+ tuple_sum: summing over the entries of a yielded tuple should give this number
+ allowed_values: optional sequence of positive integers that are considered as tuple entries
+ zero has to be in the set of allowed values
+ if None, all possible positive integers are allowed
+ ordered: if True, only tuples are returned where the entries are descending, i.e. where the entries are ordered
+
+ Examples:
+ Generate all 2-tuples where the sum of entries is 3
+ >>> list(__fixed_sum_tuples(tuple_length=2, tuple_sum=3))
+ [(0, 3), (1, 2), (2, 1), (3, 0)]
+
+ Same with ordered tuples only
+ >>> list(__fixed_sum_tuples(tuple_length=2, tuple_sum=3, ordered=True))
+ [(2, 1), (3, 0)]
+
+ Restricting the allowed values, note that zero has to be in the allowed values!
+ >>> list(__fixed_sum_tuples(tuple_length=3, tuple_sum=4, allowed_values=(0, 1, 3)))
+ [(0, 1, 3), (0, 3, 1), (1, 0, 3), (1, 3, 0), (3, 0, 1), (3, 1, 0)]
+ """
+ if not allowed_values:
+ allowed_values = set(range(0, tuple_sum + 1))
+
+ assert 0 in allowed_values and all(i >= 0 for i in allowed_values)
+
+ def recursive_helper(current_list, position, total_sum):
+ new_position = position + 1
+ if new_position < len(current_list):
+ for i in allowed_values:
+ current_list[position] = i
+ new_sum = total_sum - i
+ if new_sum < 0:
+ continue
+ for item in recursive_helper(current_list, new_position, new_sum):
+ yield item
+ else:
+ if total_sum in allowed_values:
+ current_list[-1] = total_sum
+ if not ordered:
+ yield tuple(current_list)
+ if ordered and current_list == sorted(current_list, reverse=True):
+ yield tuple(current_list)
+
+ return recursive_helper([0] * tuple_length, 0, tuple_sum)
\ No newline at end of file
diff --git a/session.py b/session.py
index 2bb6a6f0bbfa4d6b3984decbbb558d634c1c1012..259855dbb7121eac277fb998d7561c679a435de6 100644
--- a/session.py
+++ b/session.py
@@ -1,6 +1,5 @@
import sympy as sp
import numpy as np
-from pystencils.jupytersetup import *
from lbmpy.scenarios import *
from lbmpy.creationfunctions import *
from pystencils import makeSlice, showCode