lattice.py

import numpy as np
import sympy as sp
import random
from sympy import oo
from .utils.mylog import logger
from .exceptions import (
    OrderNotImplementedException,
    ImpossibleVelocityShellException,
    ReducibleShellException,
    NoSolutionException,
    InfiniteSolutionsException,
    UninitializedAttributeException,
)
from .utils.utils import (
    analyse_tensor_dimension,
    get_group,
    get_subshells,
    lattice_sum,
    approximate_ratio,
    timer,
    ps_installed,
    get_random_vector,
    velocities_for_shell,
)
from .shell import Shell


class Lattice:
    BY_NAME = {
        "D2Q9": {"dimension": 2, "order": 4, "shell_list": [1, 2, 4], "seed": 44},
        "D2V17": {
            "dimension": 2,
            "order": 6,
            "shell_list": [1, 2, 4, 8, 9],
            "seed": 44,
        },
        "D2V37": {
            "dimension": 2,
            "order": 8,
            "shell_list": [1, 2, 4, 5, 8, 9, 10, 16],
            "seed": 44,
        },
        "D3Q19": {"dimension": 3, "order": 4, "shell_list": [1, 2, 4], "seed": 44},
        "D3Q15": {"dimension": 3, "order": 4, "shell_list": [1, 3, 4], "seed": 44},
        "D3Q41-ZOT": {
            "dimension": 3,
            "order": 6,
            "shell_list": [1, 2, 3, 9, 16, 27],
            "boundary": "sup",
            "unwanted_subshells": ["221", "511"],
            "seed": 44,
        },
    }

    def __init__(
        self,
        dimension=2,
        order=4,
        shell_list=[1, 2, 4],
        seed=None,
        boundary=None,
        unwanted_subshells=[],
        svd_tolerance=1e-8,
    ):
        """Create an instance of the Lattice class without calculating the weights.
        Default values are chosen to lead to the standard D2Q9 lattice.
        Not all input values are valid. :func:`Lattice.init_by_order`

        Args:
            name: Initialize Lattice with the correct values for some known velocities. Possible names are
                D2Q9, D2V17, D2V37, D3Q17,
            dimension (int): Dimension of the resulting lattice. Default 2.
            order (int): Order of tensor isotropy. Only odd orders are
            shell_list (list): List of the Squared lengths of the velocities.
            unwanted_subshells (list) : List of strings that correspond to the type of a subshell that isn't supposed to
                to be part of lattice. Some shells are ambiguous and in the default case all possiblities are taken,
                but some lattice, like the D3Q41-ZOT need certain subshells to be remove. E.g. 2 Dimensions and
                shell 25 is added. Bother type = "50" and type = "43" could be used. Add one of them to the list of
                unwanted subshells to remove them.
            seed (float): Seed for the random number generator.
            svd_tolerance (float): tolerance to set the singular values to zero after svd.

        """
        self._dimension = dimension
        self._order = order
        self._shell_list = shell_list
        self._seed = seed
        self._svd_tolerance = svd_tolerance
        random.seed(self._seed)

        # Preperation
        self._tensor_space_dimension = 0
        self._all_velocity_vectors = []
        self._possible_tensors = []
        self._tensor_dimensions = []
        self._shells = 0
        self._unwanted_subshells = unwanted_subshells

        # LES
        self._lhs = []
        self._rhs = []

        # SVD
        self._singular_values = []
        self._solution = []

        # Polynomials
        self._coefficients = []
        self._weight_polynomials = []
        self._interval = None
        self._boundary = boundary if boundary == "sup" else "inf"

        # Final values
        self._discrete_velocities = []
        self._weights = []
        self._c_s_sq = None

    def __str__(self):
        string = "D{} - Order: {} - Shells: {}".format(
            self._dimension, self._order, str(self._shell_list)
        )
        if self._boundary:
            string += " - Boundary: {}".format(self._boundary)
        if self._unwanted_subshells:
            string += " - Unwanted Subshells: {}".format(
                "".join([s + ", " for s in self._unwanted_subshells])
            )
        return string

    @classmethod
    def from_name(cls, name):
        """Initiate the Lattice with the correct values corresponding to a know Lattice in literature.
        A complete list can be seen in Lattice.BY_NAME.keys()"""
        kwargs = cls.BY_NAME.get(name)
        if not kwargs:
            raise ValueError("No Lattice with this name is known or implemented")
        return cls(**kwargs)

    @classmethod
    def from_order(cls, dimension, order):
        """
        :param dimension:
        :param order:
        :return:
        """
        if order % 2 != 0:
            order -= 1

        for key, value in cls.BY_NAME.items():
            if value.get("dimension") == dimension and value.get("order") == order:
                return cls(**cls.BY_NAME.get(key))

        raise ValueError(
            "No Lattice known with order {} and dimension {}.".format(order, dimension)
        )
        return cls(dimension=dimension, order=order, shell_list=shell_list, seed=seed)

    @property
    def tensor_space_dimension(self):
        return np.array([len(x) for x in self._possible_tensors]).sum()

    @property
    def shells(self):
        return self._shells

    @property
    def velocities(self):
        """Get a list of all types of velocities. Velocities with zero weight are included
        For D2O9 this equals [00, 10, 11, 20]"""
        if not self._discrete_velocities and self._shells:
            self._discrete_velocities = [shell.type for shell in self._shells]
        return self._discrete_velocities

    @property
    def weights(self):
        if self._weights:
            return self._weights
        return self.calculate_weights()

    @property
    def c_s_sq(self):
        if self._c_s_sq:
            return self._c_s_sq
        return self.velocity_set()[0]

    @property
    @ps_installed
    def assignment_collection(self):
        print("okay")
        return

        if self._dimension == 3:
            return self._assignment_collection3d()
        Q = len(self.weights)
        weights = self.weights
        c_s_sq = float(self.c_s_sq)
        c_ix = self.velocities[0]
        c_iy = self.velocities[1]
        src, dst = ps.fields("src, dst({Q}): double[2D]".format(Q=Q))
        density, u_x, u_y, omega = sp.symbols("density, u_x, u_y, omega")

        density_formula = sum([src[0, 0](i) for i in range(Q)])
        vel_x_formula = (sum([src[0, 0](i) * c_ix[i] for i in range(Q)])) / density
        vel_y_formula = (sum([src[0, 0](i) * c_iy[i] for i in range(Q)])) / density

        symbolic_description = [ps.Assignment(density, density_formula),
                                ps.Assignment(u_x, vel_x_formula),
                                ps.Assignment(u_y, vel_y_formula), ]

        for i in range(Q):
            feq_formula = weights[i] * density * (
                    1
                    + (u_x * c_ix[i] + u_y * c_iy[i]) / c_s_sq
                    + (u_x * c_ix[i] + u_y * c_iy[i]) ** 2 / (2 * c_s_sq ** 2)
                    - (u_x ** 2 + u_y ** 2) / (2 * c_s_sq)
            )
            if len(c_ix) >= 17:
                feq_formula += weights[i] * density * (
                        -2 * (1 / (2 * c_s_sq)) ** 2 * (
                            c_ix[i] * u_x ** 3 + c_ix[i] * u_x * u_y ** 2 + c_iy[i] * u_y * u_x ** 2 + c_iy[
                        i] * u_y ** 3)
                        + 4 / 3 * (1 / (2 * c_s_sq)) ** 3 * (u_x * c_ix[i] + u_y * c_iy[i]) ** 3
                )
            if len(c_ix) >= 37:
                feq_formula += weights[i] * density * \
                               (1 / (6 * c_s_sq ** 2) * (
                                       (u_x * c_ix[i] + u_y * c_iy[i]) ** 4 / (4 * c_s_sq ** 2) - (
                                           3 * (u_x ** 2 + u_y ** 2) *
                                           (u_x * c_ix[i] + u_y * c_iy[i]) ** 2) / (2 * c_s_sq) + 3 * (
                                                   u_x ** 4 + u_y ** 4)))

        for i in range(Q):
            symbolic_description.append(
                ps.Assignment(dst[-1 * c_iy[i], c_ix[i]](i), omega * feq_formula + (1 - omega) * src[0, 0](i))
            )
        return ps.AssignmentCollection(symbolic_description)

    def _assignment_collection3d(self):
        pass

    def shell_from_type(self, type):
        for shell in self._shells:
            if shell.type == type:
                return shell

    def _calculate_velocity_vectors(self):
        velocity_vectors = []

        group = get_group(self._dimension)
        for shell in range(len(self._shell_list)):
            velocities = velocities_for_shell(
                dimension=self._dimension, shell=self._shell_list[shell]
            )
            if len(velocities) == 0:
                raise ImpossibleVelocityShellException(
                    "No velocity sets found for velocity shell {}.".format(shell)
                )

            # If e.g. dimension = 2, then shell 25 in ambiguous. Both shell (3,4) and (5,0) are added
            subshells = get_subshells(velocities, group)
            if len(subshells) > 1:
                # 3 4, 0 5 can only be given together
                velocities = []
                for i_subs, subshell in enumerate(subshells):
                    # get type for subshell. Example shell: all vectors of type [0, -1, 0] --> 100
                    type = "".join(
                        [str(int(x)) for x in sorted(abs(subshell[0]), reverse=True)]
                    )
                    if type not in self._unwanted_subshells:
                        velocities.append(subshell)
            else:
                velocities = [velocities]
            velocity_vectors.extend(velocities)

        self._shells = [Shell([np.array([0] * self._dimension)])]
        for final_velocities in velocity_vectors:
            self._shells.append(Shell(final_velocities))
        self._all_velocity_vectors = velocity_vectors
        return velocity_vectors

    def _fill_lhs(self):
        lhs = []
        for i, rank in enumerate(np.arange(2, self._order + 2, 2)):
            tensor_space_dimension = self._tensor_dimensions[i]
            for j in range(tensor_space_dimension):
                vector = get_random_vector(self._dimension)
                # vector = 2 * np.random.rand(self._dimension) - 1
                # vector = vector / np.linalg.norm(vector)

                rows = []
                for shell_velocities in self._all_velocity_vectors:
                    shell_sum = lattice_sum(vector, shell_velocities, rank)
                    rows.append(shell_sum)
                lhs.append(rows)
        self._lhs = np.array(lhs)
        return np.array(lhs)

    def _fill_rhs(self):
        rhs = []
        cols = self._order // 2
        for k, rank in enumerate(np.arange(2, self._order + 2, 2)):
            for j in range(self._tensor_dimensions[k]):
                rows = []
                for i in range(cols):
                    c_s_power = 2 * i + 2
                    rows.append(1 if c_s_power == rank else 0)
                rhs.append(rows)
        self._rhs = np.array(rhs)
        return np.array(rhs)

    def _svd(self):
        U, s, V = np.linalg.svd(self._lhs)
        rows = self._lhs.shape[0]
        cols = self._lhs.shape[1]

        self._singular_values = len(s)
        s = np.array([sv if sv > self._svd_tolerance else 0 for sv in s])
        rank = np.count_nonzero(s)
        rhs = np.dot(np.transpose(U), self._rhs)
        if rank < rows:
            overlap = rows - rank
            if np.linalg.norm(rhs[-overlap:]) > self._svd_tolerance:
                # NO Solution
                raise NoSolutionException(
                    "Add shells or lower the order"
                )  # change shells ?
            rows = rank
            s.resize(rows)
            rhs.resize((rows, rhs.shape[1]), refcheck=False)
            self._singular_values = len(s)

        reduced_rhs = np.zeros((rows, rhs.shape[1]))
        for i, SingularValue in enumerate(s):
            for j in range(rhs.shape[1]):
                reduced_rhs[i, j] = rhs[i, j] / SingularValue

        reduced_rhs = np.zeros((rows, rhs.shape[1]))
        for i, singular_value in enumerate(s):
            reduced_rhs[i, :] = rhs[i, :] / singular_value

        if cols - rows > 0:
            raise InfiniteSolutionsException(
                "Reduce order, remove shells, or use continue.py in "
                "https://github.com/BDuenweg/Lattice-Boltzmann-weights"
            )
        self._solution = np.dot(np.transpose(V), reduced_rhs)
        return self._solution

    def _calculate_coefficients(self):
        """Use _all_velocity_vectors and _solution to calculate the coefficients of the
        weigth polynomials. Writes them into _coefficients
        """

        solution_columns = self._order // 2
        coefficients = np.zeros((1 + solution_columns))
        coefficients[0] = 1
        for j in range(solution_columns):
            tmp = 0
            for i in range(len(self._all_velocity_vectors)):
                tmp -= self._solution[i, j] * len(self._all_velocity_vectors[i])
            coefficients[j + 1] = tmp

        coeffs = np.zeros((self._solution.shape[0] + 1, self._solution.shape[1] + 1))
        coeffs[0, :] = coefficients[-1::-1]
        coeffs[1:, :-1] = self._solution[:, -1::-1]
        self._coefficients = coeffs
        return self._coefficients

    def _calculate_valid_interval(self):
        if not self._weight_polynomials:
            raise UninitializedAttributeException(
                "Weight Polynomials not give. Calculate them first!"
            )

        interval = sp.Interval(-oo, oo)
        for equation in self._weight_polynomials:
            roots = equation.real_roots()
            if not roots:
                continue
            roots = [roots[0] - 1] + roots + [roots[-1] + 1]
            cur_interval = sp.EmptySet()
            for i in range(len(roots) - 1):
                mesh_point = (roots[i + 1] + roots[i]) / 2
                if equation.eval(mesh_point) < 0:
                    continue
                # positive weights
                if i == 0:
                    cur_interval = sp.Union(
                        cur_interval, sp.Interval(-oo, roots[i + 1])
                    )
                elif i == (len(roots) - 2):
                    cur_interval = sp.Union(cur_interval, sp.Interval(roots[i], oo))
                else:
                    cur_interval = sp.Union(
                        cur_interval, sp.Interval(roots[i], roots[i + 1])
                    )
            interval = interval.intersect(cur_interval)
        self._interval = sp.Complement(interval, sp.FiniteSet(0))
        return self._interval

    def solution_expected(self):
        """Evaluate if the order, dimension and amount of shells entered are expected to retrieve a solution.
        This is in no way a guarantee on the existence of a solution or whether or not the program
        will finish successful, but it can be used as a quick hint on the existence of a solution.

        Changes the state of self._solution and self._velocity vectors for future calculations."""

        if self._order % 2 == 1:
            self._order -= 1
            logger.warning("Only odd order valid. Decrease by one")

        for k in np.arange(2, self._order + 2, 2):
            possible_tensors = analyse_tensor_dimension(k)
            self._possible_tensors.append(possible_tensors)
        self._tensor_dimensions = [len(x) for x in self._possible_tensors]

        self._calculate_velocity_vectors()
        self._fill_lhs()
        self._fill_rhs()

        self._svd()

        return self._solution is not None

    def calculate_polynomials(self, approximate=True):
        """Main algorithm as proposed by D. Spiller and B Duenweg.

        Divided into the part where the existence of a solution and calculated together with all the velocity
        vectors. If one is expected the coefficients are calculated, the interval evaluated and the weight
        polynomials returned.

        :param approximate: Use python Fraction to approximate the calculated polynomial coefficients as rational
         number. This should give better results, since they should have a rational representation.

        :return: List of Sympy Polynomials."""

        self.solution_expected()
        self._calculate_coefficients()

        c_s_sq = sp.Symbol("c_s_sq")
        if approximate:
            self._weight_polynomials = [
                sp.Poly([approximate_ratio(x) for x in pol_coffs], c_s_sq)
                for pol_coffs in self._coefficients
            ]
        else:
            self._weight_polynomials = [
                sp.Poly(pol_coffs, c_s_sq) for pol_coffs in self._coefficients
            ]

        self._interval = self._calculate_valid_interval()
        return [] if self._interval == sp.EmptySet else self._weight_polynomials

    def calculate_weights(self, c_s_sq=None, boundary=None):
        """
        Calculate the weights for this lattice for a given speed of sound value c_s_sq, or for
        a given boundary.
        It is first checked, whether a c_s_sq value is given and then boundary.
        If both are not giving, self._boundary is taken which is defined in the
        init function.

        If no value is given, the infimum is used, i.e. the smallest possible value.

        :param c_s_sq: Float value for the speed of sound, has to be inside of the valid interval
        :param boundary: Get reduced model by using the infimum or supremum of the interval. Fallback value if
            no c_s_sq is given
        :return: weights With zero weight(s) included
        """

        if not self._weight_polynomials:
            self.calculate_polynomials()
        if not self._interval:  # Empty Interval
            return []
        if not c_s_sq:
            if not boundary:
                boundary = self._boundary
            if boundary == "inf":
                c_s_sq = self._interval.inf
            else:
                c_s_sq = self._interval.sup

        if self._interval.contains(c_s_sq) is sp.EmptySet:  # Invalid c_s_sq input
            return []

        weights = [x.eval(c_s_sq) for x in self._weight_polynomials]
        weights = [x if abs(x) >= 1e-10 else 0 for x in weights]
        for i, shell in enumerate(self.shells):
            shell.set_weight(c_s_sq, weights[i])
        self._c_s_sq = c_s_sq
        self._weights = weights
        return weights

    def velocity_set(self):
        """
        Removes unwanted shells, i.e. weight == = and
        returns the values that are important for the Lattice Boltzmann Model, i.e.
        a list of the velocity values their corresponding weights and the value for the speed of sound.

        :return: (c_s_sq, weights, velocities)
        """
        if not self.weights:
            self.calculate_weights()

        weights = []
        for shell in self.shells:
            if shell.weight == 0:
                continue
            weights += [shell.weight] * shell.size
        velocities = np.zeros((self._dimension, len(weights)), dtype=np.int32)
        i = 0
        for shell in self.shells:
            if shell.weight == 0:
                continue
            for velocity in shell.velocities:
                velocities[:, i] = velocity
                i += 1

        return self._c_s_sq, weights, velocities