Chamel/equinox api #87

cmhamel · 2024-06-28T15:00:35Z

Attempt at "re-writing part of the interface as equinox modules. This buys us a few things

Allows for jittable containers for Functionspaces, Mechanics, etc.
Easier more python like coding style with class structures, methods, and inheritance
Reduces code bloat by removing a functional style that jax enforces
Probably some other things I haven't thought of yet.

I've mainly ported QuadratureRule, FunctionSpace, and Mechanics so far.

Also I created a new class called "Domain" which removes a good bit of boiler plate code to run an optimism problem.

cmhamel · 2024-06-28T19:56:10Z

The vast majority of file changes in this PR are handled in my other PR related to unit testing and code coverage. This was branched off of that branch

codecov · 2024-06-28T23:59:22Z

Codecov Report

Attention: Patch coverage is 79.11548% with 85 lines in your changes missing coverage. Please review.

Project coverage is 76.37%. Comparing base (6782079) to head (4454bac).

Files with missing lines	Patch %	Lines
optimism/Mechanics.py	83.03%	38 Missing ⚠️
optimism/Problem.py	41.46%	24 Missing ⚠️
optimism/Domain.py	63.93%	22 Missing ⚠️
optimism/Mesh.py	96.15%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #87      +/-   ##
==========================================
- Coverage   76.40%   76.37%   -0.03%     
==========================================
  Files          61       64       +3     
  Lines        5051     5227     +176     
==========================================
+ Hits         3859     3992     +133     
- Misses       1192     1235      +43

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

cmhamel · 2024-07-08T15:40:22Z

This is a start. Mechanics.py could be cleaned up a bit more but the main question is should we make multiblock the standard rather than having two different implementations floating around.

btalamini · 2024-08-07T17:44:25Z

This is a start. Mechanics.py could be cleaned up a bit more but the main question is should we make multiblock the standard rather than having two different implementations floating around.

I’m looking forward to test driving this!

I’m also bugged about the code duplication with the multi-block stuff. I was considering going the other way and removing the multi-block version. There are usually only a few materials in our simulations. It might be easier to just have the user sum the contributions in their energy function.

ralberd · 2024-08-29T20:53:33Z

This is a start. Mechanics.py could be cleaned up a bit more but the main question is should we make multiblock the standard rather than having two different implementations floating around.

I’m looking forward to test driving this!

I’m also bugged about the code duplication with the multi-block stuff. I was considering going the other way and removing the multi-block version. There are usually only a few materials in our simulations. It might be easier to just have the user sum the contributions in their energy function.

I agree with Brandon, it may be cleaner to just have the single block mechanics functions and then the Domain class that you added can handle looping over blocks to compute energies, etc. I think it makes sense for Domain to contain the knowledge of blocks.

ralberd

This looks pretty good. Once you address the comments and the TODO items it should be good to go.

ralberd · 2024-08-29T20:13:39Z

optimism/Domain.py

+  def get_bc_values(self, U):
+    return self.dof.get_bc_values(U)


Is this function necessary? It looks like get_ubcs covers this functionality.

ralberd · 2024-08-29T20:15:17Z

optimism/Domain.py

+
+  def update_field(self, Uu, p):
+    Ubc = self.get_ubcs(p)
+    return self.create_field(Uu, Ubc)


Can this just be return self.dof.create_field(Uu, Ubc)? Then you could get rid of self.create_field

ralberd · 2024-08-29T20:18:01Z

optimism/Domain.py

+  def create_field(self, Uu, Ubc):
+    return self.dof.create_field(Uu, Ubc)
+
+  def create_unknowns(self):


Suggested change

def create_unknowns(self):

def initialize_unknowns(self):

ralberd · 2024-08-29T20:20:26Z

optimism/Domain.py

+    Uu = self.dof.get_unknown_values(np.zeros(self.mesh.coords.shape))
+    return Uu
+
+  def energy_function(self, Uu, p):


Is there a way for this to be extended easily to incorporate things like traction energy? Or would a different Domain-like class be required for each energy type?

I think this is the one place where we'll need to do some functional programming under the hood. We initializing the Domain we can call some called "create_energy_function" that will construct things appropriately depending upon what inputs are provided to the constructor.

ralberd · 2024-08-29T20:25:54Z

optimism/FunctionSpace.py

+    shapeGrads: Float[Array, "ne nqpe npe nd"]
+    mesh: any
+    quadratureRule: QuadratureRule.QuadratureRule
+    isAxisymmetric: bool


Is the idea to eventually move the contents of construct_function_space into the constructor for this class?

Yeah exactly. I was going to wait to do that until a future PR since, otherwise it would've changed a large number of test files.

ralberd · 2024-08-29T21:08:54Z

optimism/Mesh.py

@@ -64,8 +83,8 @@ def construct_structured_mesh(Nx, Ny, xExtent, yExtent, elementOrder=1, useBubbl
    return mesh


-def get_blocks(mesh, blockNames):
-    return tuple(mesh.blocks[name] for name in blockNames)
+# def get_blocks(mesh, blockNames):


Is this no longer used?

optimism/Objective.py

optimism/Problem.py

ralberd · 2024-08-29T21:21:12Z

@btalamini @tupek2. Once this gets merged in, what should be the plan to deal with the old functions? How important is backwards compatibility?

…the functionspace level and created a class structure for mechanics that we can now derive off while jitting and vmapping and gradding.

…k class could be removed eventually. Had to switch from jit to filter_jit in a few places in tests, but minor changes.

…greatly pairs down the effort to setup a new simulation.

…minor change to the json mesh reader and a test.

Co-authored-by: ralberd <[email protected]>

cmhamel marked this pull request as draft June 28, 2024 17:47

cmhamel marked this pull request as ready for review June 29, 2024 00:49

cmhamel changed the title ~~[WIP] Chamel/equinox api [DO NOT MERGE]~~ Chamel/equinox api Jun 30, 2024

cmhamel requested review from ralberd, tupek2 and btalamini July 8, 2024 15:39

ralberd reviewed Aug 29, 2024

View reviewed changes

cmhamel and others added 17 commits September 11, 2024 17:25

adding equinox as dependency.

4d3d84c

a lot of changes to allow for jittable almost everything. Started at …

12c31a9

…the functionspace level and created a class structure for mechanics that we can now derive off while jitting and vmapping and gradding.

some cleanup of dynamics.

ca472f6

finished mapping over Mechanics.py to an equinox interface. Multibloc…

422d80a

…k class could be removed eventually. Had to switch from jit to filter_jit in a few places in tests, but minor changes.

adding example with the new equinox api. This is just a start but it …

646acfd

…greatly pairs down the effort to setup a new simulation.

some modifications to make multi block the standard. This required a …

cba834a

…minor change to the json mesh reader and a test.

making dynamics kinetic energy support multiple blocks.

bfe37bc

patched up dynamics methods for hessian

b6462df

modifying tests to remove MultiBlock from name

a235c3b

finished up the failing newmark tests and think this looks good for now.

a764461

moving the problem class into its own file

449cd25

some more testing on Domain.py.

428dd8a

Remove comment

5e70069

Co-authored-by: ralberd <[email protected]>

Removing comment

9b42ba3

Co-authored-by: ralberd <[email protected]>

Adding a missing jit

560ff7d

Co-authored-by: ralberd <[email protected]>

Removing comment

2422f79

Co-authored-by: ralberd <[email protected]>

Adding missing jit

613775b

Co-authored-by: ralberd <[email protected]>

cmhamel and others added 6 commits September 11, 2024 17:25

Adding missing jit

946acd8

Co-authored-by: ralberd <[email protected]>

Adding missing jit

5be5613

Co-authored-by: ralberd <[email protected]>

Adding missing jit

30c332b

Co-authored-by: ralberd <[email protected]>

Removing comment

816bce1

Co-authored-by: ralberd <[email protected]>

Method rename

78520dc

Co-authored-by: ralberd <[email protected]>

Method rename

4454bac

Co-authored-by: ralberd <[email protected]>

cmhamel force-pushed the chamel/equinox-api branch from 158daef to 4454bac Compare September 11, 2024 23:25

cmhamel mentioned this pull request Oct 25, 2024

moving QuadratureRule over to an equinox.Module and patched up a few … #96

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Chamel/equinox api #87

Chamel/equinox api #87

cmhamel commented Jun 28, 2024

cmhamel commented Jun 28, 2024

codecov bot commented Jun 28, 2024 •

edited

Loading

cmhamel commented Jul 8, 2024

btalamini commented Aug 7, 2024

ralberd commented Aug 29, 2024

ralberd left a comment

ralberd Aug 29, 2024

ralberd Aug 29, 2024

ralberd Aug 29, 2024

ralberd Aug 29, 2024

cmhamel Sep 3, 2024

ralberd Aug 29, 2024

cmhamel Sep 3, 2024

ralberd Aug 29, 2024

ralberd commented Aug 29, 2024

Chamel/equinox api #87

Are you sure you want to change the base?

Chamel/equinox api #87

Conversation

cmhamel commented Jun 28, 2024

cmhamel commented Jun 28, 2024

codecov bot commented Jun 28, 2024 • edited Loading

Codecov Report

cmhamel commented Jul 8, 2024

btalamini commented Aug 7, 2024

ralberd commented Aug 29, 2024

ralberd left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ralberd commented Aug 29, 2024

codecov bot commented Jun 28, 2024 •

edited

Loading