Update loops to use jax calls and JIT some functions #27
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
SachinJalan
reviewed
Aug 13, 2023
|
|
||
| Parameters | ||
| ---------- | ||
| obj : diffmpm.node.Nodes, diffmpm.particle._ParticlesState |
There was a problem hiding this comment.
I think it should be vel in the docstring or obj in the function, same issue with all the functions.
SachinJalan
reviewed
Aug 13, 2023
| Damping forces on the nodes. | ||
| loc : ArrayLike | ||
| Locations of all the nodes. Expected shape (nnodes, 1, ndim) | ||
| initialized: bool |
There was a problem hiding this comment.
There are no parameters initialized,data in this method.
SachinJalan
reviewed
Aug 13, 2023
diffmpm/particle.py
Outdated
Comment on lines
57
to
64
| initialized: bool | ||
| `False` if particle property arrays like mass need to be initialized. | ||
| If `True`, they are set to values from `data`. | ||
| data: tuple | ||
| Tuple of length 13 that sets arrays for mass, density, volume, | ||
| velocity, acceleration, momentum, strain, stress, strain_rate, | ||
| dstrain, f_ext, reference_loc and volumetric_strain_centroid. |
There was a problem hiding this comment.
These parameters are not present in this method also docstring for init_vel can be added.
SachinJalan
reviewed
Aug 13, 2023
Comment on lines
+146
to
+158
| Parameters | ||
| ---------- | ||
| state: | ||
| Current state | ||
| elements: diffmpm._ElementState | ||
| Elements that the particles are present in, and are used to | ||
| compute the particles' volumes. | ||
| total_elements: int | ||
| Total elements present in `elements`. | ||
| """ |
There was a problem hiding this comment.
docstring not matching function arguments
SachinJalan
reviewed
Aug 13, 2023
| initialized: Optional[bool] = None, | ||
| volume: Optional[ArrayLike] = None, | ||
| ) -> Quad4NState: | ||
| """Initialize Linear1D. |
There was a problem hiding this comment.
Suggested change
| """Initialize Linear1D. | |
| """Initialize Quad4N. |
SachinJalan
reviewed
Aug 13, 2023
|
|
||
| def init_state( | ||
| self, | ||
| nelements: int, |
There was a problem hiding this comment.
Suggested change
| nelements: int, | |
| nelements: chex.ArrayDevice, |
SachinJalan
reviewed
Aug 13, 2023
| self, | ||
| nelements: int, | ||
| total_elements: int, | ||
| el_len: float, |
There was a problem hiding this comment.
Suggested change
| el_len: float, | |
| el_len: chex.ArrayDevice, |
To make the imports easier from the materials subdir, also restructured other files. This moves `MPM` to a separate file so as to remove circular imports for materials module.
chahak13
force-pushed
the
jax_apply_loops
branch
from
a83eaeb
to
fcec796
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
No description provided.