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=[],
        eq_order=None,
        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.
            eq_order (int): Explicitly specify the order of accuracy for the equilibrium distribution function.
                Only relevant if AssignmentCollections are used. For 2D resonable defaults are used.
            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._eq_order = eq_order
        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._q = None
        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 q(self):
        if self._q:
            return self._q
        self._q = len(self.velocity_set()[1])
        return self._q

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

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

    @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 eq_order(self):
        """Calculate the eq order based on the amount of velocities given and dimension.
        According to Philippi et. al, 17 and 37 velocities are the minimum amount respectively
        to use for third and fourth order accurate models. """
        if self._eq_order:
            return self._eq_order

        if self._dimension == 2:
            if self.q >= 37:
                self._eq_order = 4
            elif self.q >= 17:
                self._eq_order = 3
            else:
                self._eq_order = 2
        if self._dimension == 3:
            # Should be done explicitly for three dimensions using the constructor
            self._eq_order = 2
        return self._eq_order

    @property
    @ps_installed
    def assignment_collection(self, ps):
        if self._dimension == 3:
            return self._assignment_collection3d()
        c_s_sq, weights, velocities = self.velocity_set()
        Q = self.q
        c_ix = velocities[0]
        c_iy = velocities[1]
        src, dst = ps.fields("src({Q}), 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[0, 0] * (
                    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 self.eq_order >= 3:
                feq_formula += weights[i] * density[0, 0] * (
                    -1 / (2 * c_s_sq**2) * (u_x * c_ix[i] + u_y * c_iy[i]) * (u_x ** 2 + u_y ** 2)
                    + 4 / 3 * (1 / (2 * c_s_sq)) ** 3 * (u_x * c_ix[i] + u_y * c_iy[i]) ** 3
                )
            if self.eq_order >= 4:
                feq_formula += weights[i] * density[0, 0] * \
                               (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)))

            symbolic_description.append(
                ps.Assignment(dst[-1 * c_iy[i], c_ix[i]](i), omega * feq_formula + (1 - omega) * src[0, 0](i))
            )
            print("okakychanp")
        return ps.AssignmentCollection(symbolic_description)

    @property
    @ps_installed
    def _assignment_collection3d(self, ps):
        c_s_sq, weights, velocities = self.velocity_set()
        Q = self.q
        c_ix = velocities[0]
        c_iy = velocities[1]
        c_iz = velocities[2]
        src, dst = ps.fields("src({Q}), dst({Q}): double[2D]".format(Q=Q))
        density, u_x, u_y, u_z, omega = sp.symbols("density, u_x, u_y, u_z, 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
        vel_z_formula = (sum([src[0, 0](i) * c_iz[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),
                                ps.Assignment(u_z, vel_z_formula),]

        for i in range(Q):
            feq_formula = weights[i] * density[0, 0] * (
                    1
                    + (u_x * c_ix[i] + u_y * c_iy[i] + u_z * c_iz[i]) / c_s_sq
                    + (u_x * c_ix[i] + u_y * c_iy[i] + u_z * c_iz[i]) ** 2 / (2 * c_s_sq ** 2)
                    - (u_x ** 2 + u_y ** 2 + u_z ** 2) / (2 * c_s_sq)
            )
            if self.eq_order >= 3:
                feq_formula += weights[i] * density[0, 0] * (
                        -1 / (2 * c_s_sq ** 2) * (u_x * c_ix[i] + u_y * c_iy[i] + u_z * c_iz[i]) * (u_x ** 2 + u_y ** 2 + u_z ** 2)
                        + 4 / 3 * (1 / (2 * c_s_sq)) ** 3 * (u_x * c_ix[i] + u_y * c_iy[i] + u_z * c_iz[i]) ** 3
                )
            if self.eq_order >= 4:
                feq_formula += weights[i] * density[0, 0] * \
                               (1 / (6 * c_s_sq ** 2) * (
                                       (u_x * c_ix[i] + u_y * c_iy[i] + u_z * c_iz[i]) ** 4 / (4 * c_s_sq ** 2) - (
                                           3 * (u_x ** 2 + u_y ** 2 + u_z ** 2) *
                                           (u_x * c_ix[i] + u_y * c_iy[i] + u_z * c_iz[i]) ** 2) / (2 * c_s_sq) + 3 * (
                                                   u_x ** 4 + u_y ** 4 + u_z ** 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 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