from fourier import fourier
from bte     import bte 
import matinverse as mi
from jax import numpy as jnp

grid    = 40  #Grid resolution
n_phi   = 48  #Grid resolution
L       = 100 #Simulation Domain nm
R       = 30  #Radious of the conic filter in nm
Knt     = 1   #Knudsen number (referred to the whole cell)


#f is the fourier solver and b is the bte one. A generic cost function can be developed as in the example below

f    = mi.compose(fourier(grid=grid,directions=[[1,0],[0,1]],eta_fourier=1))
b    = mi.compose(bte(grid=grid,n_phi=n_phi,Knt=1,directions=[[1,0],[0,1]]))

kd = [0.2,0.5]
def objective(x):

        (kx,ky),aux  = f(x)

        g  = jnp.sqrt((kx-kd[0])**2 + (ky-kd[1])**2)

        return g,((kx,ky),aux)


#As an option, this can take "inequality_constraints=func" where func is a function defined exactly like "objective" above
x = mi.run(objective,grid = grid,L = L,R = R,n_betas=10,maxiter=30,dim=2,min_porosity=0.2)