#!/usr/bin/env python
# coding: utf-8

# In[1]:


from lbmpy.session import *


# # Demo: Create lbmpy Method from Scratch
# 
# <img src='../img/collision_space.svg' width="90%">
# 
# 
# ### Defining transformation to collision space

# In[2]:


from lbmpy.moments import moment_matrix, moments_up_to_component_order, exponents_to_polynomial_representations
moment_exponents = list(moments_up_to_component_order(2, 2))
moment_exponents


# In[3]:


moments = exponents_to_polynomial_representations(moment_exponents)
moments


# In[4]:


d2q9 = LBStencil(Stencil.D2Q9, ordering='walberla')
d2q9.plot()


# In[5]:


M = moment_matrix(moments, stencil=d2q9)
M


# ### Defining the Equilibrium Distribution, and Computation of Conserved Quantities

# In[6]:


from lbmpy.equilibrium import ContinuousHydrodynamicMaxwellian
from lbmpy.methods import DensityVelocityComputation

equilibrium = ContinuousHydrodynamicMaxwellian(dim=2, compressible=True, order=2)
cqc = DensityVelocityComputation(d2q9, True, False)

tuple(zip(moments, equilibrium.moments(moments)))


# ### Defining Relaxation Behaviour and Creating the Method

# In[7]:


from lbmpy.methods.creationfunctions import create_from_equilibrium

omega = sp.symbols("omega")
relaxation_rate_dict = {moment : omega for moment in moments}

force_model = forcemodels.Guo(sp.symbols("F_:2"))
method = create_from_equilibrium(d2q9, equilibrium, cqc, relaxation_rate_dict)
method


# ### Example of a update equation without simplifications

# In[8]:


collision_rule = method.get_collision_rule()
collision_rule


# ### Generic simplification strategy - common subexpresssion elimination

# In[9]:


generic_strategy = ps.simp.SimplificationStrategy()
generic_strategy.add(ps.simp.sympy_cse)
generic_strategy.create_simplification_report(collision_rule)


# ### A custom simplification strategy for moment-based methods

# In[10]:


simplification_strategy = create_simplification_strategy(method)
simplification_strategy.create_simplification_report(collision_rule)
simplification_strategy.add(ps.simp.sympy_cse)