docstrings added access modifiers to functions

Author: SachinJalan

diffmpm/material.py

@@ -154,28 +154,28 @@ def __init__(self, material_properties):
         ---------
         material_properties: dict
             Dictionary with material properties. For Bingham
-        material, 'density','youngs_modulus','poisson_ratio','tau0','mu'
-        and 'critical_shear_rate' are required keys.
-
-        ndim: int
-            Dimension of the problem supports 1D and 2D
+        material, 'density','youngs_modulus','poisson_ratio','tau0','mu',
+        'ndim and 'critical_shear_rate' are required keys.
+
+        Methods
+        -------
+        initialise_state_variables
+            Initialises the state variables for the Bingham material
+        compute_stress
+            computes the stress for the Bingham material particles
+        
         """
         self.validate_props(material_properties)
         self.ndim = material_properties["ndim"]
-        density = material_properties["density"]
         youngs_modulus = material_properties["youngs_modulus"]
         poisson_ratio = material_properties["poisson_ratio"]
-        tau0 = material_properties["tau0"]
-        mu_ = material_properties["mu"]
-        critical_shear_rate = material_properties["critical_shear_rate"]
+        self.state_variables=["pressure"]
         # Calculate the bulk modulus
         bulk_modulus = youngs_modulus / (3.0 * (1.0 - 2.0 * poisson_ratio))
         compressibility_multiplier_ = 1.0
         # Special Material Properties
-        if "incompressible" in material_properties:
-            incompressible = material_properties["incompressible"]
-            if incompressible:
-                compressibility_multiplier_ = 0.0
+        if material_properties.get("incompressible", False):
+            compressibility_multiplier_ = 0.0
         self.properties = {
             **material_properties,
             "bulk_modulus": bulk_modulus,
@@ -191,16 +191,30 @@ def initialise_state_variables(particles):
         state_vars["pressure"] = jnp.zeros((particles.loc.shape[0]))
         return state_vars
 
-    # State Variables
-    def state_variables():
-        return ["pressure"]
-
     # Compute the pressure
-    def thermodynamic_pressure(self, volumetric_strain):
+    def __thermodynamic_pressure(self, volumetric_strain):
         return -self.properties["bulk_modulus"] * volumetric_strain
 
     # Compute the stress
-    def compute_stress(self, dstrain, particles, state_vars):
+    def compute_stress(self, dstrain, particles, state_vars:dict):
+        """
+        Computes the stress for the Bingham material
+
+        Arguments
+        ---------
+        dstrain: array_like
+            The strain rate tensor for the particles
+        particles: diffmpm.particles.Particles
+        state_vars: dict {str: jnp.ndarray}
+            dictionary containig the string as the name of the
+            property and the jnp.ndarray shape (nparticles, 1) as the
+            values of the property at each particle.
+             
+        Returns
+        -------
+        updated_stress: jnp.ndarray
+            The updated stress for the particles expected shape (nparticles, 6,1)
+        """
         shear_rate_threshold = 1e-15
         # dirac delta in Voigt notation
         dirac_delta = jnp.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0]).reshape((6, 1))
@@ -218,7 +232,7 @@ def compute_stress(self, dstrain, particles, state_vars):
         )
 
         @jit
-        def compute_stress_per_particle(
+        def __compute_stress_per_particle(
             particle_strain_rate,
             self,
             state_vars_pressure,
@@ -228,21 +242,20 @@ def compute_stress_per_particle(
             strain_r = particle_strain_rate
             # Convert strain rate to rate of deformation tensor
             strain_r = strain_r.at[-3:].multiply(0.5)
-            """
-            Rate of shear = sqrt(2 * D_ij * D_ij)
-            Since D (D_ij) is in Voigt notation (D_i), and the definition above is in
-            matrix, the last 3 components have to be doubled D_ij * D_ij = D_0^2 +
-            D_1^2 + D_2^2 + 2*D_3^2 + 2*D_4^2 + 2*D_5^2 Yielding is defined: rate of
-            shear > critical_shear_rate_^2 Checking yielding from strain rate vs
-            critical yielding shear rate
-            """
+
+            # Rate of shear = sqrt(2 * Strain_r_ij * Strain_r_ij)
+            # Strain_r is in Voigt notation so the above formula reduces in the voigt notation to
+            # sqrt(Strain_r_0^2 + Strain_r_1^2 + Strain_r_2^2 + 2*Strain_r_3^2 
+            # + 2*Strain_r_4^2 + 2*Strain_r_5^2)
+            # When shear rate> critical_shear_rate^2 then the material is yielding
+
             shear_rate = jnp.sqrt(
                 2.0 * (strain_r.T @ (strain_r) + strain_r[-3:].T @ strain_r[-3:])
             )
-            """
-            Apparent_viscosity maps shear rate to shear stress
-            Check if shear rate is 0
-            """
+
+            # Apparent_viscosity maps shear rate to shear stress
+            # Check if shear rate is 0
+
             apparent_viscosity_true = 2.0 * (
                 (self.properties["tau0"] / shear_rate[0, 0]) + self.properties["mu"]
             )
@@ -251,17 +264,15 @@ def compute_stress_per_particle(
                 * self.properties["critical_shear_rate"]
             )
             apparent_viscosity = lax.select(condition, apparent_viscosity_true, 0.0)
-            """
-            Compute shear change to volumetric
-            tau deviatoric part of cauchy stress tensor
-            """
+
+            # Compute volumetric tau
+
             tau = apparent_viscosity * strain_r
-            """
-            von Mises criterion
-            trace of second invariant J2 of deviatoric stress in matrix form
-            Since tau is in Voigt notation, only the first three numbers matter
-            yield condition trace of the invariant > tau0^2
-            """
+            # von Mises criterion
+            # yield condition trace of the invariant > tau0^2 
+            # and trace can be found using the first 3 numbers of tau 
+            # as tau is in voigt notation
+            
             trace_invariant = 0.5 * jnp.dot(tau[:3, 0], tau[:3, 0])
             tau = lax.cond(
                 trace_invariant < (self.properties["tau0"] * self.properties["tau0"]),
@@ -272,23 +283,26 @@ def compute_stress_per_particle(
             # update pressure
             state_vars_pressure += self.properties[
                 "compressibility_multiplier"
-            ] * self.thermodynamic_pressure(dvolumetric_strain_per_particle)
-            """
-            Update volumetric and deviatoric stress
-            thermodynamic pressure is from material point
-            stress = -thermodynamic_pressure I + tau, where I is identity matrix or
-            direc_delta in Voigt notation
-            """
+            ] * self.__thermodynamic_pressure(dvolumetric_strain_per_particle)
+
+            # Update volumetric and deviatoric stress
+            # thermodynamic pressure is from material point
+            # stress = -thermodynamic_pressure I + tau, where I is identity matrix or
+            # direc_delta in Voigt notation
+
             updated_stress_per_particle = (
                 -(state_vars_pressure)
                 * dirac_delta
                 * self.properties["compressibility_multiplier"]
                 + tau
             )
             return updated_stress_per_particle, state_vars_pressure
-
+        
+        # using vmap to vectorise the function compute stress per particle
+        # for all the particles using the first dimension of the strain rate
+        # and the dvolumetric_strain and state_vars pressure
         updated_stress, state_vars["pressure"] = vmap(
-            compute_stress_per_particle, in_axes=(0, None, 0, 0, None)
+            __compute_stress_per_particle, in_axes=(0, None, 0, 0, None)
         )(
             particles.strain_rate,
             self,