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


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 = []

        # Output
        self._coefficients = []
        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

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

        self._velocities = velocities_amount
        self._shells = len(total_subshells) + 1

        logger.info("Velocities: {} Shells: {}".format(self._velocities, self._shells))

        for i, shell in enumerate(velocity_vectors):
            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))
        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

    @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 self._interval

    @timer
    def calculate_weights(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]

        self._interval = 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 self._weight_polynomials]
        weights = [x if abs(x) >= 1e-10 else 0 for x in weights]
        return weights