lattice.py

import numpy as np
import sympy as sp
from sympy import oo
from lbmweights.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, fill_rhs, approximate_ratio, timer, get_random_vector


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

        # LES
        self._tensor_space_dimension = 0
        self._possible_tensors = []
        self._tensor_dimensions = []
        self._velocities = 0
        self._shells = 0

        # SVD
        self._amount_singular_values = 0
        self._singular_values = []

        self._weight_polynomials = []
        self._interval = 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):
        self._shells

    @property
    def velocities(self):
        self._velocities

    @staticmethod
    def velocities_for_shell(dimension, shell):
        """
        :param dimension: Lattice dimension
        :param shell: Squared velocity of the shell
        :return: list of velocities with squared length equal to shell
        """
        velocities = []
        linear_lattice_size = 2 * shell + 1
        full_lattice_size = linear_lattice_size ** dimension
        for site in range(full_lattice_size):
            tmp = site
            cur_vel_squared = 0
            cur_vec = []

            for dim in range(dimension):
                coordinate = tmp % linear_lattice_size
                tmp = (tmp - coordinate) // linear_lattice_size
                shifted_coordinate = coordinate - shell
                cur_vec.append(shifted_coordinate)
                cur_vel_squared += shifted_coordinate ** 2

            if cur_vel_squared == shell:
                velocities.append(np.array(cur_vec, dtype=int))

        return velocities

    @timer
    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 interval

    @timer
    def calculate_weights(self):
        logger.debug("Entering Calculate Weights")
        logger.info(str(self))
        import random
        random.seed(self._seed)
        print("My Seed: {}".format(self._seed))
        if self._order % 2 == 1:
            self._order -= 1
            logger.warning("Odd order for isotropy entered. This package only handles symmetric lattices, therefore"
                           "the equations for odd orders are automatically fulfilled. "
                           "Reduction of order to {}".format(self._order))

        # Do tensor dimension analysis for every order up to the desired
        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]

        # FROM HERE  ...
        velocities_amount = 1
        grand_total_list = []
        total_subshells = []

        group = get_group(self._dimension)
        for shell in range(len(self._shell_list)):
            velocities = self.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))

            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]
            grand_total_list.extend(velocities)
            velocities_amount += len(velocities)

        self._group  = group
        self._grand_total_list = grand_total_list
        self._velocities = velocities_amount
        self._shells = len(total_subshells) + 1
        logger.info("Velocities: {} Shells: {}".format(self._velocities, self._shells))

        for i, shell in enumerate(grand_total_list):
            number_of_velocities = len(shell)
            type = tuple(np.sort(abs(shell[0])))
            #logger.info("Shell number {} with c_i^2 = {} and {} velocities of type {}".format(i+1, (shell[0]**2).sum(), number_of_velocities, type))

        # TO HERE: Only Calculate GrandTotalList

        # Here: ...
        # SpacialDimension = self._dimension
        # MaxTensorRank = self._order
        # ListOfTensordimensions = [len(x) for x in self._possible_tensors]
        # GrandTotalList = grand_total_list == list shells which are list of arrays with the actual values
        # Arguments = <class 'dict'>: {'d': 2, 'm': 4, 'c': [1, 2, 4], 's': 44, 'y': 'test': 'quiet': 'write_latex':}

        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 grand_total_list:
                    shell_sum = lattice_sum(vector, shell_velocities, rank)
                    rows.append(shell_sum)
                lhs.append(rows)
        self._lhs = np.array(lhs)
        A = np.array(lhs)
        b = fill_rhs(self._order, self._tensor_dimensions)
        self._rhs = np.copy(b)

        # --------------------------------------------------------------------------------------------------

        U, s, V = np.linalg.svd(A)
        self._U = np.copy(U)
        self._s = np.copy(s)
        self._V = np.copy(V)
        rows = A.shape[0]
        cols = A.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), b)
        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

        x = np.copy(reduced_rhs)
        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")

        # ----------------------------------------------- Unique solution
        solution = np.dot(np.transpose(V), reduced_rhs)
        self._solution = solution
        solution_columns = self._order // 2
        assert solution_columns == solution.shape[1]

        coefficients = np.zeros((1 + solution_columns))
        coefficients[0] = 1
        for j in range(solution_columns):
            tmp = 0
            for i in range(len(grand_total_list)):
                tmp -= solution[i, j] * len(grand_total_list[i])
            coefficients[j+1] = tmp

        self._coefficients = np.array(coefficients)
        # -------------------------------------------------- Prettify for sympy
        weight_polynomials = []
        c_s_sq = sp.Symbol("c_s_sq")
        coeffs = np.zeros((solution.shape[0] + 1, solution.shape[1] + 1))
        coeffs[0, :] = coefficients[-1::-1]
        coeffs[1:, :-1] = solution[:,-1::-1]

        apporx_coeffs = [[approximate_ratio(x) for x in pol_coffs] for pol_coffs in coeffs]

        weight_polynomials = [sp.Poly([approximate_ratio(x) for x in pol_coffs], c_s_sq) for pol_coffs in coeffs]
        # -------------------------------------------------- Calculate Roots

        self._weight_polynomials = weight_polynomials
        self.calculate_valid_interval()
        logger.info(self._interval.boundary)
        if self._interval == sp.EmptySet():
            return []
        weights = [sp.N(x.eval(self._interval.boundary.inf)) for x in weight_polynomials]
        weights = [x if abs(x) >= 1e-10 else 0 for x in weights]
        return weights