Add type annotations, docstrings

benchmarks/2d/uniaxial_nodal_forces/test_benchmark.py

@@ -1,6 +1,8 @@
 import os
 from pathlib import Path
+
 import jax.numpy as jnp
+
 from diffmpm import MPM
 
 

benchmarks/2d/uniaxial_particle_traction/test_benchmark.py

@@ -1,6 +1,8 @@
 import os
 from pathlib import Path
+
 import jax.numpy as jnp
+
 from diffmpm import MPM
 
 

benchmarks/2d/uniaxial_stress/test_benchmark.py

@@ -1,6 +1,8 @@
 import os
 from pathlib import Path
+
 import jax.numpy as jnp
+
 from diffmpm import MPM
 
 

diffmpm/__init__.py

@@ -40,7 +40,7 @@ def __init__(self, filepath):
             raise ValueError("Wrong type of solver specified.")
 
     def solve(self):
-        """Solve the MPM simulation."""
+        """Solve the MPM simulation using JIT solver."""
         arrays = self.solver.solve_jit(
             self._config.parsed_config["external_loading"]["gravity"],
         )

diffmpm/cli/mpm.py

@@ -3,11 +3,12 @@
 from diffmpm import MPM
 
 
-@click.command()
+@click.command()  # type: ignore
 @click.option(
     "-f", "--file", "filepath", required=True, type=str, help="Input TOML file"
 )
 @click.version_option(package_name="diffmpm")
 def mpm(filepath):
+    """CLI utility for MPM."""
     solver = MPM(filepath)
     solver.solve()

diffmpm/constraint.py

@@ -3,7 +3,18 @@
 
 @register_pytree_node_class
 class Constraint:
-    def __init__(self, dir, velocity):
+    """Generic velocity constraints to apply on nodes or particles."""
+
+    def __init__(self, dir: int, velocity: float):
+        """Contains 2 govering parameters.
+
+        Attributes
+        ----------
+        dir : int
+            Direction in which constraint is applied.
+        velocity : float
+            Constrained velocity to be applied.
+        """
         self.dir = dir
         self.velocity = velocity
 
@@ -16,16 +27,15 @@ def tree_unflatten(cls, aux_data, children):
         return cls(*aux_data)
 
     def apply(self, obj, ids):
-        """
-        Apply constraint values to the passed object.
+        """Apply constraint values to the passed object.
 
-        Arguments
-        ---------
+        Parameters
+        ----------
         obj : diffmpm.node.Nodes, diffmpm.particle.Particles
             Object on which the constraint is applied
         ids : array_like
             The indices of the container `obj` on which the constraint
-        will be applied.
+            will be applied.
         """
         obj.velocity = obj.velocity.at[ids, :, self.dir].set(self.velocity)
         obj.momentum = obj.momentum.at[ids, :, self.dir].set(