Merge pull request #1 from chahak13/jaxport

[WIP] Port library to JAX.

Author: kks32

README.md

@@ -0,0 +1 @@
+# DiffMPM

diffmpm/__init__.py

@@ -0,0 +1 @@
+#!/usr/bin/env python3

diffmpm/material.py

@@ -1,9 +1,33 @@
-import jax.numpy as jnp
+from jax.tree_util import register_pytree_node_class
 
-class LinearElastic:
-    def  __init__(self, E, density):
+
+@register_pytree_node_class
+class Material:
+    """
+    Base material class.
+    """
+
+    def __init__(self, E, density):
+        """
+        Initialize material properties.
+
+        Arguments
+        ---------
+        E : float
+            Young's modulus of the material.
+        density : float
+            Density of the material.
+        """
         self.E = E
         self.density = density
 
-    def update_stress(self, particle, dt):
-        particle.stress+=particle.dstrain*self.E
+    def tree_flatten(self):
+        return (tuple(), (self.E, self.density))
+
+    @classmethod
+    def tree_unflatten(cls, aux_data, children):
+        del children
+        return cls(*aux_data)
+
+    def __repr__(self):
+        return f"Material(E={self.E}, density={self.density})"