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, get_random_vector, velocities_for_shell
from .shell import Shell


class Lattice:

    def __init__(self, dimension=2, order=4, shell_list=[1, 2, 4], seed=None, svd_tolerance=1e-8):
        """
        :param dimension:
        :param order:
        :param shell_list:
        :param seed:
        :param tolerance: 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._velocities = 0
        self._shells = 0

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

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

        # Polynomials
        self._coefficients = []
        self._weight_polynomials = []
        self._interval = None

        # 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._seed:
            string += " - Seed: {}".format(self._seed)
        return string

    @classmethod
    def init_by_order(cls, dimension, order, seed=None):
        """
        :param dimension:
        :param order:
        :return:
        """
        shell_list = []
        if dimension == 2:
            if order == 4:
                shell_list = [1, 2, 4]
            elif order == 6:
                shell_list = [1, 2, 4, 8, 9]
            elif order == 8:
                shell_list = [1, 2, 4, 5, 8, 9, 10, 16]

        elif dimension == 3:
            if order == 4:
                shell_list = [1, 2, 4]

        if not shell_list:
            raise OrderNotImplementedException(
                "Cannot initialize by {} for dimension {}. Lower the order or try to come up with the shells yourself.".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):
        return self._discrete_velocities

    def calculate_velocity_vectors(self):
        velocities_amount = 1
        velocity_vectors = []
        total_subshells = []

        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)
            total_subshells.append(len(subshells))
            if len(subshells) > 1:
                # 3 4, 0 5 can only be given together
                for i_subs, subshell in enumerate(subshells):
                    type = tuple(np.sort(abs(subshell[0])))
                velocities = subshells
            else:
                velocities = [velocities]
            velocity_vectors.extend(velocities)
            velocities_amount += len(velocities)

        logger.info("Velocities: {} Shells: {}".format(self._velocities, self._shells))
        self._velocities = velocities_amount
        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:
            # TODO Better
            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):
        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 calculate_polynomials(self):
        c_s_sq = sp.Symbol("c_s_sq")
        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()

        self.calculate_coefficients()

        apporx_coeffs = [[approximate_ratio(x) for x in pol_coffs] for pol_coffs in self._coefficients]

        self._weight_polynomials = [sp.Poly([approximate_ratio(x) for x in pol_coffs], c_s_sq) for pol_coffs in self._coefficients]
        if not self._weight_polynomials:
            raise NoSolutionException("Something went wrong")

        self._interval = self.calculate_valid_interval()
        logger.info(self._interval.boundary)
        if self._interval == sp.EmptySet():
            return []
        return self._weight_polynomials

    def calculate_reduced_weights(self, boundary="inf"):
        """
        Calculate the weights of this lattice for a reduced model.
        If the smallest or biggest value of the valid interval is taken, they reduce
        the amount of shells necessary by one.
        Overwrites the weight values in self.shell

        :param boundary: "inf" --> smallest possible value is taken, otherwise the largest
        :return: list of weights in the order of their shells.
        """
        return self.calculate_weights(boundary=boundary)

    def calculate_weights(self, c_s_sq=None, boundary="inf"):
        """
        Calculate the weights for this lattice for a given speed of sound value c_s_sq.
        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:
        """
        if not self._weight_polynomials:
            self.calculate_polynomials()
        if not self._interval:  # Empty Interval
            return []
        if not c_s_sq:
            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 = [sp.N(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):
        """
        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): c_s_sq
        """

        weights = []
        for shell in self.shells:
            if shell.weight == 0:
                continue
            weights += [shell.weight] * shell.size
        velocities = np.zeros((self._dimension, len(weights)))
        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