Update loops to use jax calls and JIT some functions - #27

Draft

chahak13 wants to merge 21 commits into

geoelements:mainfrom

chahak13:jax_apply_loops

Commits on Aug 17, 2023

Remove extra example files
chahak13
committed
0880975
Restructure code and add materials subdir
chahak13
committed
3e309e0
Pass `particles` as arg to compute stress
chahak13
committed
8b4f1c0
Add Newtonian material
chahak13
committed
763af13
Convert list loops to tree_map calls and jit funcs
chahak13
committed
6120e44
Run benchmark on cpu
chahak13
committed
0055007
Temp commit: To be squashed
chahak13
committed
f285088
Make Nodes as state
chahak13
committed
3518368
Working nonfrozen node state
chahak13
committed
25289e7
Working frozen Nodestate
chahak13
committed
00f9822
Working nonfrozen particle state
chahak13
committed
28e7bff
Working frozen particle state
chahak13
committed
4158e11
Reset node state in scheme and separate bc apply
chahak13
committed
3d0e43e
Return changed element arrays, not new state
chahak13
committed
f6e1880
Working on completely stateless design
chahak13
committed
5f3b5f7
Working uniaxial stress and nodal forces benchmark
chahak13
committed
09bd261
Broken traction
chahak13
committed
a07cadf
Update tests
chahak13
committed
b4b2e38
Change loops to scan and add optim example
chahak13
committed
1034d4f
Update optim file
chahak13
committed
03f112c
Add uniaxial stress optim example plot
chahak13
committed
fcec796