Package: crikit
CRIKit: The Constitutive Relation Inference ToolkitQuick Start|Documentation|Installation GuideOverviewCRIKit integratesFEniCSandPyadjointwith machine learning libraries
likeJAXandTensorFlow, and provides tools to infer physically-compatible constitutive relations from sparse, noisy observations of a system modeled by partial differential equations. CRIKit bridges the FEniCS world with those of
JAX and TensorFlow by storingcovering mapsbetween abstractSpaceclasses that
represent spaces like a FEniCSFunctionSpaceor a space of JAX arrays of a
particular shape, or adirect sumof multipleSpaces.CRIKit also provides tools to help perform post-processing, such as observation operators, as well as a collection of loss functions.Quick StartConstructing And Optimizing a CRThis guide will show you the basics of constructing and optimizing a simple CR
that represents linear elasticity, assuming that you're already familiar with
the basics ofFEniCS. You can compare the mechanics of CRIKit to that of FEniCS
directly by comparing this example tothe 2D linear elasticity examplefromNumerical tours of Computational Mechanics using FEniCS. The primary difference
between the model shown here and the linked example in the previous sentence is that
here we use a geometrically nonlinear model, as described in the documentation for
thelibCEED hyperelasticity example.fromcrikitimport*importjaxfromjaximportnumpyasjnpimportnumpyasnpfromfunctoolsimportpartial# set up mesh, FunctionSpace, etcfe_order=2dims=2Nx,Ny=50,5L=20.H=1.mesh=RectangleMesh(Point(0.,0.),Point(L,H),Nx,Ny)V=VectorFunctionSpace(mesh,"CG",fe_order)quad_params={'quadrature_degree':fe_order+1}set_default_covering_params(domain=mesh.ufl_domain(),quad_params=quad_params)u=Function(V)defleft_boundary(x,on_boundary):returnnear(x[0],0.)bcs=[DirichletBC(V,Constant((0.,0.)),left_boundary)]# these will tell CRIKit what the inputs and ouputs to the CR# are so that we can automatically generate the scalar and form-invariants# Let's suppose you want the Cauchy stress tensor as a function of the# strain sym(grad(u))input_types=(TensorType.make_symmetric(2,dims,'strain'),)output_type=TensorType.make_symmetric(2,dims,'stress')# initial guess of parametersYoungs=1.0e5Poisson=0.3lmbda=(Youngs*Poisson)/((1+Poisson)*(1-2*Poisson))mu=Youngs/(2*(1+Poisson))# since this is 2-d, we need to use a modified version of lambda# to make our initial guesses physicallmbda=2*lmbda*mu/(lmbda+2*mu)theta=array([lmbda,mu])defcr_func(invariants,theta):lmbda,mu=thetareturnjnp.array([lmbda*jnp.log1p(invariants[0]),2*mu])cr=CR(output_type,input_types,cr_func,params=[theta])# If you're in a Jupyter notebook, run this at the bottom of a cell instead of# calling `print()` on it to get neatly-rendered HTML output.# This function shows you a description of the scalar and form invariants of `cr`# in the order they are placed in the arraysprint(cr.invariant_descriptions())# set the default covering params for crikit.covering so we can automatically# generate covering maps between spaces of FEniCS Functions and JAX arrays# Let's just pretend that degree 3 is sufficient quadrature for whatever problem# we're solvingquad_params={'quadrature_degree':3}set_default_covering_params(domain=mesh.ufl_domain(),quad_params=quad_params)# create_ufl_standins() returns a tuple of objects that can act as standins# for the output of a CR. You can't directly call the CR on the inputs because# the CR expects JAX arrays as an input, not a FEniCS Function. You'll instead have# to assembly the variational form F using assemble_with_cr(), which will generate# a covering map from the space of FEniCS Functions to the space of JAX arrays# using crikit.covering (and likewise from the output JAX array space to a space of# `crikit.fe.Function`s), use it to get appropriate arguments, call the CR, and project the result# back into a Functiontarget_shape=tuple(iforiincr.target.shape()ifi!=-1)standin_sigma,=create_ufl_standins((target_shape,))# create your form as if standin_sigma were (cr(sym(grad(u)))v=TestFunction(V)# external forcef=Constant((0,-1e-3),name='force')F=inner(standin_sigma,sym(grad(v)))*dx-inner(f,v)*dx# define a new sub-tape that records the actions of this equationwithpush_tape():# a function that we can assemble the variational form into# using the `tensor` kwarg of `crikit.assemble()`, which# is directly passed on to `crikit.fe.backend.assemble()`# (e.g. `fenics.assemble()`)residual=Function(V)# input to the CR is sym(grad(u))assemble_with_cr(F,cr,sym(grad(u)),standin_sigma,tensor=residual,quad_params=quad_params)ucontrol=Control(u)# a ReducedFunction to represent the residual as a function of `u`res_rf=ReducedFunction(residual,ucontrol)# an object to represent the equation defined abovered_eq=ReducedEquation(res_rf,bcs,homogenize_bcs(bcs))# and an object to solve it. Make sure your .petscrc is set appropriately!# if you want to pass an assembled Jacobian, use 'jmat_type' : 'assembled',# but if you want the solver to instead use the matrix-free Jacobian action,# pass 'jmat_type' : 'action'solver=SNESSolver(red_eq,{'jmat_type':'assembled'})pred_u=solver.solve(ucontrol)# define a loss function and an observernum_slices=100seed=0# sliced quadratic Wasserstein distanceloss=SlicedWassersteinDistance(V,num_slices,jax.random.PRNGKey(seed),p=2)classObservedSubDomain(SubDomain):definside(self,x,on_boundary):...# return appropriate True/False if x is in the observed subdomain or not# observe only on a given SubDomainobserver=SubdomainObserver(mesh,ObservedSubDomain())# get your observations from somewhere as a Function in Vobs=...err=loss(observer(obs),observer(pred_u))Jhat=ReducedFunctional(err,Control(theta))#check the derivativeh=np.random.randn(*theta.shape)v=array(1.0)# test the adjointasserttaylor_test(Jhat,theta,h,v=v)>=1.9# choose an optimization methodopt_method='L-BFGS-B'optimal_params=minimize(Jhat,method=opt_method)InstallationFirst, installFEniCS. Then you can install the latest development
version of this package by runningpip install .or, to install in editable mode (useful for devs),pip install -e .There are four optional sets of dependencies that can be installed by listing
any of them (with no spaces) in square brackets. Or if you want to install all
of them, you can use the keyall, so the two lines below are equivalent.pip install -e .[test,visualization,tensorflow,doc]
pip install -e .[all]If you would rather install the latest release version of CRIKit, you can
instead runpip install crikitMake sure you have install CRIKit into an environment with a working FEniCS
installation, or the build might fail. In particular, if you encounter
errors buildingpetsc4py, ensure that the FEniCS installation in your
environment works.Setting Up a Conda EnvironmentFEniCS provides a conda package, so installation into a conda environment is simple.condacreate--namefenics2019python=3.7--no-default-packages-y-q
condaactivatefenics2019

condainstall-cconda-forgefenics=2019.1.0-y-q
pipinstall-e.[all]Whenever you want to enter this environment, runconda activate fenics2019.DocumentationDocumentation is done withSphinx.Documentation can be built by runningmakein thedocsfolder. By default,
that will create HTML documentation indocs/build/html, and you can view it by
opening updocs/build/html/index.htmlin your favorite browser.Runmake todo=trueto include Todo blocks in the output. Otherwise, they won't show up.You can also runmake helpto see a list of other formats that can be built. For instance,make coveragecreates a file showing what classes/functions are missing documentation.make doctesttests the example code in the documentation.ExamplesExample programs are in theexamplesfolder. Run with the--helpargument to
see command-line parameters. For example,cd examples/p-stokes
python p-stokes.py --helpTestsYou can run the main tests with the command below.python3 -m pytest testsYou can run a specific test by runningpython3 -m pytest tests -k test_name,
wheretest_nameis all or part of the test name. For example,python3 -m pytest tests/crikit/test_assemble_with_cr.py -k test_assemble_with_cr_scalarorpython3 -m pytest tests/crikit/test_assemble_with_cr.py -k scalarSome implementation files and documentation files have doctests. Those can be run like so:python -m pytest --doctest-modules --rootdir tests crikit
cd docs && make doctestDeveloper style guidelinesCRIKit uses the auto-formatterblackto ensure the code has a consistent style. To automatically
run it before each git commit, use the following commands.pip install -e .[dev]
pre-commit installSupportThis material is based upon work supported by the National Science Foundation under Grant No. 1835825 and 1835792.
Any opinions, findings, and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the National Science Foundation.