From 0880975d3c533abe8698150600a73d6f3816d974 Mon Sep 17 00:00:00 2001 From: Chahak Mehta Date: Sat, 15 Jul 2023 13:22:31 -0700 Subject: [PATCH 01/21] Remove extra example files --- examples/mpm-nodal-forces.toml | 73 --------------------- examples/mpm-uniaxial-stress.toml | 61 ------------------ examples/particles-2d-nodal-force.json | 74 ---------------------- examples/particles-2d-uniaxial-stress.json | 26 -------- 4 files changed, 234 deletions(-) delete mode 100644 examples/mpm-nodal-forces.toml delete mode 100644 examples/mpm-uniaxial-stress.toml delete mode 100644 examples/particles-2d-nodal-force.json delete mode 100644 examples/particles-2d-uniaxial-stress.json diff --git a/examples/mpm-nodal-forces.toml b/examples/mpm-nodal-forces.toml deleted file mode 100644 index cf01f1e..0000000 --- a/examples/mpm-nodal-forces.toml +++ /dev/null @@ -1,73 +0,0 @@ -# The `meta` group contains top level attributes that govern the -# behaviour of the MPM Solver. -# -# Attributes: -# title: The title of the experiment. This is just for the user's -# reference. -# type: The type of simulation to be used. Allowed values are -# {"MPMExplicit"} -# scheme: The MPM Scheme used for simulation. Allowed values are -# {"usl", "usf"} -# dt: Timestep used in the simulation. -# nsteps: Number of steps to run the simulation for. -[meta] -title = "uniaxial-nodal-traction" -type = "MPMExplicit" -dimension = 2 -scheme = "usf" -dt = 0.001 -nsteps = 301 -velocity_update = true - -[output] -type = "hdf5" -file = "results/example_2d_out.hdf5" -step_frequency = 5 - -[mesh] -# type = "file" -# file = "mesh-1d.txt" -# boundary_nodes = "boundary-1d.txt" -# particle_element_ids = "particles-elements.txt" -type = "generator" -nelements = [3, 1] -element_length = [0.1, 0.1] -particle_element_ids = [0] -element = "Quadrilateral4Node" - -[[mesh.constraints]] -node_ids = [0, 4] -dir = 0 -velocity = 0.0 - -[[materials]] -id = 0 -density = 1000 -poisson_ratio = 0 -youngs_modulus = 1000000 -type = "LinearElastic" - -[[particles]] -file = "examples/particles-2d-nodal-force.json" -material_id = 0 -init_velocity = 0.0 - -[external_loading] -gravity = [0, 0] - -[[external_loading.concentrated_nodal_forces]] -node_ids = [3, 7] -math_function_id = 0 -dir = 0 -force = 0.05 - -[[external_loading.particle_surface_traction]] -pset = [1] -dir = 1 -math_function_id = 0 -traction = 10.5 - -[[math_functions]] -type = "Linear" -xvalues = [0.0, 0.5, 1.0] -fxvalues = [0.0, 1.0, 1.0] diff --git a/examples/mpm-uniaxial-stress.toml b/examples/mpm-uniaxial-stress.toml deleted file mode 100644 index 4f8065e..0000000 --- a/examples/mpm-uniaxial-stress.toml +++ /dev/null @@ -1,61 +0,0 @@ -# The `meta` group contains top level attributes that govern the -# behaviour of the MPM Solver. -# -# Attributes: -# title: The title of the experiment. This is just for the user's -# reference. -# type: The type of simulation to be used. Allowed values are -# {"MPMExplicit"} -# scheme: The MPM Scheme used for simulation. Allowed values are -# {"usl", "usf"} -# dt: Timestep used in the simulation. -# nsteps: Number of steps to run the simulation for. -[meta] -title = "uniaxial-stress" -type = "MPMExplicit" -dimension = 2 -scheme = "usf" -dt = 0.01 -nsteps = 10 -velocity_update = false - -[output] -format = "npz" -folder = "results/" -step_frequency = 5 - -[mesh] -# type = "file" -# file = "mesh-1d.txt" -# boundary_nodes = "boundary-1d.txt" -# particle_element_ids = "particles-elements.txt" -type = "generator" -nelements = [1, 1] -element_length = [1, 1] -particle_element_ids = [0] -element = "Quadrilateral4Node" - -[[mesh.constraints]] -node_ids = [0, 1] -dir = 1 -velocity = 0.0 - -[[mesh.constraints]] -node_ids = [2, 3] -dir = 1 -velocity = -0.01 - -[[materials]] -id = 0 -density = 1 -poisson_ratio = 0 -youngs_modulus = 1000 -type = "LinearElastic" - -[[particles]] -file = "examples/particles-2d-uniaxial-stress.json" -material_id = 0 -init_velocity = [1.0, 0.0] - -[external_loading] -gravity = [0, 0] diff --git a/examples/particles-2d-nodal-force.json b/examples/particles-2d-nodal-force.json deleted file mode 100644 index f0143b8..0000000 --- a/examples/particles-2d-nodal-force.json +++ /dev/null @@ -1,74 +0,0 @@ -[ - [ - [ - 0.025, - 0.025 - ] - ], - [ - [ - 0.075, - 0.025 - ] - ], - [ - [ - 0.125, - 0.025 - ] - ], - [ - [ - 0.175, - 0.025 - ] - ], - [ - [ - 0.225, - 0.025 - ] - ], - [ - [ - 0.275, - 0.025 - ] - ], - [ - [ - 0.025, - 0.075 - ] - ], - [ - [ - 0.075, - 0.075 - ] - ], - [ - [ - 0.125, - 0.075 - ] - ], - [ - [ - 0.175, - 0.075 - ] - ], - [ - [ - 0.225, - 0.075 - ] - ], - [ - [ - 0.275, - 0.075 - ] - ] -] \ No newline at end of file diff --git a/examples/particles-2d-uniaxial-stress.json b/examples/particles-2d-uniaxial-stress.json deleted file mode 100644 index 3b22d51..0000000 --- a/examples/particles-2d-uniaxial-stress.json +++ /dev/null @@ -1,26 +0,0 @@ -[ - [ - [ - 0.25, - 0.25 - ] - ], - [ - [ - 0.75, - 0.25 - ] - ], - [ - [ - 0.75, - 0.75 - ] - ], - [ - [ - 0.25, - 0.75 - ] - ] -] \ No newline at end of file From 3e309e09cf23a5b23fec8526085cb02a083a0cdf Mon Sep 17 00:00:00 2001 From: Chahak Mehta Date: Sat, 15 Jul 2023 14:54:19 -0700 Subject: [PATCH 02/21] Restructure code and add materials subdir 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. --- .../uniaxial_nodal_forces/test_benchmark.py | 2 +- .../test_benchmark.py | 2 +- .../2d/uniaxial_stress/test_benchmark.py | 2 +- diffmpm/__init__.py | 44 +------------ diffmpm/cli/mpm.py | 2 +- diffmpm/io.py | 3 +- diffmpm/materials/__init__.py | 3 + diffmpm/materials/_base.py | 48 ++++++++++++++ .../linear_elastic.py} | 66 +------------------ diffmpm/materials/newtonian.py | 1 + diffmpm/materials/simple.py | 18 +++++ diffmpm/mpm.py | 42 ++++++++++++ diffmpm/particle.py | 6 +- tests/test_element.py | 2 +- tests/test_material.py | 3 +- tests/test_particle.py | 2 +- 16 files changed, 126 insertions(+), 120 deletions(-) create mode 100644 diffmpm/materials/__init__.py create mode 100644 diffmpm/materials/_base.py rename diffmpm/{material.py => materials/linear_elastic.py} (55%) create mode 100644 diffmpm/materials/newtonian.py create mode 100644 diffmpm/materials/simple.py create mode 100644 diffmpm/mpm.py diff --git a/benchmarks/2d/uniaxial_nodal_forces/test_benchmark.py b/benchmarks/2d/uniaxial_nodal_forces/test_benchmark.py index ae72923..4dcb077 100644 --- a/benchmarks/2d/uniaxial_nodal_forces/test_benchmark.py +++ b/benchmarks/2d/uniaxial_nodal_forces/test_benchmark.py @@ -3,7 +3,7 @@ import jax.numpy as jnp -from diffmpm import MPM +from diffmpm.mpm import MPM def test_benchmarks(): diff --git a/benchmarks/2d/uniaxial_particle_traction/test_benchmark.py b/benchmarks/2d/uniaxial_particle_traction/test_benchmark.py index 356d0a3..995ca16 100644 --- a/benchmarks/2d/uniaxial_particle_traction/test_benchmark.py +++ b/benchmarks/2d/uniaxial_particle_traction/test_benchmark.py @@ -3,7 +3,7 @@ import jax.numpy as jnp -from diffmpm import MPM +from diffmpm.mpm import MPM def test_benchmarks(): diff --git a/benchmarks/2d/uniaxial_stress/test_benchmark.py b/benchmarks/2d/uniaxial_stress/test_benchmark.py index f04e820..0dd6af8 100644 --- a/benchmarks/2d/uniaxial_stress/test_benchmark.py +++ b/benchmarks/2d/uniaxial_stress/test_benchmark.py @@ -3,7 +3,7 @@ import jax.numpy as jnp -from diffmpm import MPM +from diffmpm.mpm import MPM def test_benchmarks(): diff --git a/diffmpm/__init__.py b/diffmpm/__init__.py index faa8316..a138300 100644 --- a/diffmpm/__init__.py +++ b/diffmpm/__init__.py @@ -1,47 +1,5 @@ from importlib.metadata import version -from pathlib import Path -import diffmpm.writers as writers -from diffmpm.io import Config -from diffmpm.solver import MPMExplicit - -__all__ = ["MPM", "__version__"] +__all__ = ["__version__"] __version__ = version("diffmpm") - - -class MPM: - def __init__(self, filepath): - self._config = Config(filepath) - mesh = self._config.parse() - out_dir = Path(self._config.parsed_config["output"]["folder"]).joinpath( - self._config.parsed_config["meta"]["title"], - ) - - write_format = self._config.parsed_config["output"].get("format", None) - if write_format is None or write_format.lower() == "none": - writer_func = None - elif write_format == "npz": - writer_func = writers.NPZWriter().write - else: - raise ValueError(f"Specified output format not supported: {write_format}") - - if self._config.parsed_config["meta"]["type"] == "MPMExplicit": - self.solver = MPMExplicit( - mesh, - self._config.parsed_config["meta"]["dt"], - velocity_update=self._config.parsed_config["meta"]["velocity_update"], - sim_steps=self._config.parsed_config["meta"]["nsteps"], - out_steps=self._config.parsed_config["output"]["step_frequency"], - out_dir=out_dir, - writer_func=writer_func, - ) - else: - raise ValueError("Wrong type of solver specified.") - - def solve(self): - """Solve the MPM simulation using JIT solver.""" - arrays = self.solver.solve_jit( - self._config.parsed_config["external_loading"]["gravity"], - ) - return arrays diff --git a/diffmpm/cli/mpm.py b/diffmpm/cli/mpm.py index aebc4ba..0b4b9d7 100644 --- a/diffmpm/cli/mpm.py +++ b/diffmpm/cli/mpm.py @@ -1,6 +1,6 @@ import click -from diffmpm import MPM +from diffmpm.mpm import MPM @click.command() # type: ignore diff --git a/diffmpm/io.py b/diffmpm/io.py index d6e4573..66447f5 100644 --- a/diffmpm/io.py +++ b/diffmpm/io.py @@ -1,11 +1,10 @@ import json import tomllib as tl -from collections import namedtuple import jax.numpy as jnp from diffmpm import element as mpel -from diffmpm import material as mpmat +from diffmpm import materials as mpmat from diffmpm import mesh as mpmesh from diffmpm.constraint import Constraint from diffmpm.forces import NodalForce, ParticleTraction diff --git a/diffmpm/materials/__init__.py b/diffmpm/materials/__init__.py new file mode 100644 index 0000000..ce35083 --- /dev/null +++ b/diffmpm/materials/__init__.py @@ -0,0 +1,3 @@ +from diffmpm.materials._base import _Material +from diffmpm.materials.simple import SimpleMaterial +from diffmpm.materials.linear_elastic import LinearElastic diff --git a/diffmpm/materials/_base.py b/diffmpm/materials/_base.py new file mode 100644 index 0000000..d30b15b --- /dev/null +++ b/diffmpm/materials/_base.py @@ -0,0 +1,48 @@ +import abc +from typing import Tuple + + +class _Material(abc.ABC): + """Base material class.""" + + _props: Tuple[str, ...] + + def __init__(self, material_properties): + """Initialize material properties. + + Parameters + ---------- + material_properties: dict + A key-value map for various material properties. + """ + self.properties = material_properties + + # @abc.abstractmethod + def tree_flatten(self): + """Flatten this class as PyTree Node.""" + return (tuple(), self.properties) + + # @abc.abstractmethod + @classmethod + def tree_unflatten(cls, aux_data, children): + """Unflatten this class as PyTree Node.""" + del children + return cls(aux_data) + + @abc.abstractmethod + def __repr__(self): + """Repr for Material class.""" + ... + + @abc.abstractmethod + def compute_stress(self): + """Compute stress for the material.""" + ... + + def validate_props(self, material_properties): + for key in self._props: + if key not in material_properties: + raise KeyError( + f"'{key}' should be present in `material_properties` " + f"for {self.__class__.__name__} materials." + ) diff --git a/diffmpm/material.py b/diffmpm/materials/linear_elastic.py similarity index 55% rename from diffmpm/material.py rename to diffmpm/materials/linear_elastic.py index 09230d4..098c10e 100644 --- a/diffmpm/material.py +++ b/diffmpm/materials/linear_elastic.py @@ -1,58 +1,11 @@ -import abc -from typing import Tuple - import jax.numpy as jnp from jax.tree_util import register_pytree_node_class - -class Material(abc.ABC): - """Base material class.""" - - _props: Tuple[str, ...] - - def __init__(self, material_properties): - """Initialize material properties. - - Parameters - ---------- - material_properties: dict - A key-value map for various material properties. - """ - self.properties = material_properties - - # @abc.abstractmethod - def tree_flatten(self): - """Flatten this class as PyTree Node.""" - return (tuple(), self.properties) - - # @abc.abstractmethod - @classmethod - def tree_unflatten(cls, aux_data, children): - """Unflatten this class as PyTree Node.""" - del children - return cls(aux_data) - - @abc.abstractmethod - def __repr__(self): - """Repr for Material class.""" - ... - - @abc.abstractmethod - def compute_stress(self): - """Compute stress for the material.""" - ... - - def validate_props(self, material_properties): - for key in self._props: - if key not in material_properties: - raise KeyError( - f"'{key}' should be present in `material_properties` " - f"for {self.__class__.__name__} materials." - ) +from ._base import _Material @register_pytree_node_class -class LinearElastic(Material): +class LinearElastic(_Material): """Linear Elastic Material.""" _props = ("density", "youngs_modulus", "poisson_ratio") @@ -114,18 +67,3 @@ def compute_stress(self, dstrain): """Compute material stress.""" dstress = self.de @ dstrain return dstress - - -@register_pytree_node_class -class SimpleMaterial(Material): - _props = ("E", "density") - - def __init__(self, material_properties): - self.validate_props(material_properties) - self.properties = material_properties - - def __repr__(self): - return f"SimpleMaterial(props={self.properties})" - - def compute_stress(self, dstrain): - return dstrain * self.properties["E"] diff --git a/diffmpm/materials/newtonian.py b/diffmpm/materials/newtonian.py new file mode 100644 index 0000000..e5a0d9b --- /dev/null +++ b/diffmpm/materials/newtonian.py @@ -0,0 +1 @@ +#!/usr/bin/env python3 diff --git a/diffmpm/materials/simple.py b/diffmpm/materials/simple.py new file mode 100644 index 0000000..d9cf15d --- /dev/null +++ b/diffmpm/materials/simple.py @@ -0,0 +1,18 @@ +from jax.tree_util import register_pytree_node_class + +from ._base import _Material + + +@register_pytree_node_class +class SimpleMaterial(_Material): + _props = ("E", "density") + + def __init__(self, material_properties): + self.validate_props(material_properties) + self.properties = material_properties + + def __repr__(self): + return f"SimpleMaterial(props={self.properties})" + + def compute_stress(self, dstrain): + return dstrain * self.properties["E"] diff --git a/diffmpm/mpm.py b/diffmpm/mpm.py new file mode 100644 index 0000000..b06eaee --- /dev/null +++ b/diffmpm/mpm.py @@ -0,0 +1,42 @@ +from pathlib import Path + +import diffmpm.writers as writers +from diffmpm.io import Config +from diffmpm.solver import MPMExplicit + + +class MPM: + def __init__(self, filepath): + self._config = Config(filepath) + mesh = self._config.parse() + out_dir = Path(self._config.parsed_config["output"]["folder"]).joinpath( + self._config.parsed_config["meta"]["title"], + ) + + write_format = self._config.parsed_config["output"].get("format", None) + if write_format is None or write_format.lower() == "none": + writer_func = None + elif write_format == "npz": + writer_func = writers.NPZWriter().write + else: + raise ValueError(f"Specified output format not supported: {write_format}") + + if self._config.parsed_config["meta"]["type"] == "MPMExplicit": + self.solver = MPMExplicit( + mesh, + self._config.parsed_config["meta"]["dt"], + velocity_update=self._config.parsed_config["meta"]["velocity_update"], + sim_steps=self._config.parsed_config["meta"]["nsteps"], + out_steps=self._config.parsed_config["output"]["step_frequency"], + out_dir=out_dir, + writer_func=writer_func, + ) + else: + raise ValueError("Wrong type of solver specified.") + + def solve(self): + """Solve the MPM simulation using JIT solver.""" + arrays = self.solver.solve_jit( + self._config.parsed_config["external_loading"]["gravity"], + ) + return arrays diff --git a/diffmpm/particle.py b/diffmpm/particle.py index 1bb3d70..586fec0 100644 --- a/diffmpm/particle.py +++ b/diffmpm/particle.py @@ -6,7 +6,7 @@ from jax.typing import ArrayLike from diffmpm.element import _Element -from diffmpm.material import Material +from diffmpm.materials import _Material @register_pytree_node_class @@ -16,7 +16,7 @@ class Particles(Sized): def __init__( self, loc: ArrayLike, - material: Material, + material: _Material, element_ids: ArrayLike, initialized: Optional[bool] = None, data: Optional[Tuple[ArrayLike, ...]] = None, @@ -27,7 +27,7 @@ def __init__( ---------- loc: ArrayLike Location of the particles. Expected shape (nparticles, 1, ndim) - material: diffmpm.material.Material + material: diffmpm.materials._Material Type of material for the set of particles. element_ids: ArrayLike The element ids that the particles belong to. This contains diff --git a/tests/test_element.py b/tests/test_element.py index 50881d9..ff8d92e 100644 --- a/tests/test_element.py +++ b/tests/test_element.py @@ -5,7 +5,7 @@ from diffmpm.element import Quadrilateral4Node from diffmpm.forces import NodalForce from diffmpm.functions import Unit -from diffmpm.material import SimpleMaterial +from diffmpm.materials import SimpleMaterial from diffmpm.particle import Particles diff --git a/tests/test_material.py b/tests/test_material.py index 2e041d7..66cb4dc 100644 --- a/tests/test_material.py +++ b/tests/test_material.py @@ -1,7 +1,6 @@ import jax.numpy as jnp import pytest - -from diffmpm.material import LinearElastic, SimpleMaterial +from diffmpm.materials import LinearElastic, SimpleMaterial material_dstrain_stress_targets = [ ( diff --git a/tests/test_particle.py b/tests/test_particle.py index d7dedaa..d67bc2f 100644 --- a/tests/test_particle.py +++ b/tests/test_particle.py @@ -2,7 +2,7 @@ import pytest from diffmpm.element import Quadrilateral4Node -from diffmpm.material import SimpleMaterial +from diffmpm.materials import SimpleMaterial from diffmpm.particle import Particles From 8b4f1c07777fb5ca03bb2d81931d263fd702c91d Mon Sep 17 00:00:00 2001 From: Chahak Mehta Date: Sat, 15 Jul 2023 16:58:52 -0700 Subject: [PATCH 03/21] Pass `particles` as arg to compute stress --- diffmpm/materials/_base.py | 3 ++- diffmpm/materials/linear_elastic.py | 5 ++-- diffmpm/materials/simple.py | 5 ++-- diffmpm/particle.py | 9 ++++++- tests/test_material.py | 38 +++++++++++++++++++++++------ 5 files changed, 46 insertions(+), 14 deletions(-) diff --git a/diffmpm/materials/_base.py b/diffmpm/materials/_base.py index d30b15b..896b206 100644 --- a/diffmpm/materials/_base.py +++ b/diffmpm/materials/_base.py @@ -6,6 +6,7 @@ class _Material(abc.ABC): """Base material class.""" _props: Tuple[str, ...] + properties: dict def __init__(self, material_properties): """Initialize material properties. @@ -35,7 +36,7 @@ def __repr__(self): ... @abc.abstractmethod - def compute_stress(self): + def compute_stress(self, particles): """Compute stress for the material.""" ... diff --git a/diffmpm/materials/linear_elastic.py b/diffmpm/materials/linear_elastic.py index 098c10e..5a008d4 100644 --- a/diffmpm/materials/linear_elastic.py +++ b/diffmpm/materials/linear_elastic.py @@ -9,6 +9,7 @@ class LinearElastic(_Material): """Linear Elastic Material.""" _props = ("density", "youngs_modulus", "poisson_ratio") + state_vars = () def __init__(self, material_properties): """Create a Linear Elastic material. @@ -63,7 +64,7 @@ def _compute_elastic_tensor(self): ] ) - def compute_stress(self, dstrain): + def compute_stress(self, particles): """Compute material stress.""" - dstress = self.de @ dstrain + dstress = self.de @ particles.dstrain return dstress diff --git a/diffmpm/materials/simple.py b/diffmpm/materials/simple.py index d9cf15d..77b57ca 100644 --- a/diffmpm/materials/simple.py +++ b/diffmpm/materials/simple.py @@ -6,6 +6,7 @@ @register_pytree_node_class class SimpleMaterial(_Material): _props = ("E", "density") + state_vars = () def __init__(self, material_properties): self.validate_props(material_properties) @@ -14,5 +15,5 @@ def __init__(self, material_properties): def __repr__(self): return f"SimpleMaterial(props={self.properties})" - def compute_stress(self, dstrain): - return dstrain * self.properties["E"] + def compute_stress(self, particles): + return particles.dstrain * self.properties["E"] diff --git a/diffmpm/particle.py b/diffmpm/particle.py index 586fec0..04f2581 100644 --- a/diffmpm/particle.py +++ b/diffmpm/particle.py @@ -69,6 +69,11 @@ def __init__( self.reference_loc = jnp.zeros_like(self.loc) self.dvolumetric_strain = jnp.zeros((self.loc.shape[0], 1)) self.volumetric_strain_centroid = jnp.zeros((self.loc.shape[0], 1)) + self.state_vars = {} + if self.material.state_vars: + self.state_vars = self.material.initialize_state_variables( + self.loc.shape[0] + ) else: ( self.mass, @@ -87,6 +92,7 @@ def __init__( self.reference_loc, self.dvolumetric_strain, self.volumetric_strain_centroid, + self.state_vars, ) = data # type: ignore self.initialized = True @@ -112,6 +118,7 @@ def tree_flatten(self): self.reference_loc, self.dvolumetric_strain, self.volumetric_strain_centroid, + self.state_vars, ) aux_data = (self.material,) return (children, aux_data) @@ -319,7 +326,7 @@ def compute_stress(self, *args): particles. The stress calculated by the material is then added to the particles current stress values. """ - self.stress = self.stress.at[:].add(self.material.compute_stress(self.dstrain)) + self.stress = self.stress.at[:].add(self.material.compute_stress(self)) def update_volume(self, *args): """Update volume based on central strain rate.""" diff --git a/tests/test_material.py b/tests/test_material.py index 66cb4dc..f81cfb0 100644 --- a/tests/test_material.py +++ b/tests/test_material.py @@ -1,27 +1,48 @@ import jax.numpy as jnp import pytest from diffmpm.materials import LinearElastic, SimpleMaterial +from diffmpm.particle import Particles -material_dstrain_stress_targets = [ +particles_dstrain_stress_targets = [ ( - SimpleMaterial({"E": 10, "density": 1}), + Particles( + jnp.array([[0.5, 0.5]]).reshape(1, 1, 2), + SimpleMaterial({"E": 10, "density": 1}), + jnp.array([0]), + ), jnp.ones((1, 6, 1)), jnp.ones((1, 6, 1)) * 10, ), ( - LinearElastic({"density": 1, "youngs_modulus": 10, "poisson_ratio": 1}), + Particles( + jnp.array([[0.5, 0.5]]).reshape(1, 1, 2), + LinearElastic({"density": 1, "youngs_modulus": 10, "poisson_ratio": 1}), + jnp.array([0]), + ), jnp.ones((1, 6, 1)), jnp.array([-10, -10, -10, 2.5, 2.5, 2.5]).reshape(1, 6, 1), ), ( - LinearElastic({"density": 1000, "youngs_modulus": 1e7, "poisson_ratio": 0.3}), + Particles( + jnp.array([[0.5, 0.5]]).reshape(1, 1, 2), + LinearElastic( + {"density": 1000, "youngs_modulus": 1e7, "poisson_ratio": 0.3} + ), + jnp.array([0]), + ), jnp.array([0.001, 0.0005, 0, 0, 0, 0]).reshape(1, 6, 1), jnp.array([1.63461538461538e4, 12500, 0.86538461538462e4, 0, 0, 0]).reshape( 1, 6, 1 ), ), ( - LinearElastic({"density": 1000, "youngs_modulus": 1e7, "poisson_ratio": 0.3}), + Particles( + jnp.array([[0.5, 0.5]]).reshape(1, 1, 2), + LinearElastic( + {"density": 1000, "youngs_modulus": 1e7, "poisson_ratio": 0.3} + ), + jnp.array([0]), + ), jnp.array([0.001, 0.0005, 0, 0.00001, 0, 0]).reshape(1, 6, 1), jnp.array( [1.63461538461538e4, 12500, 0.86538461538462e4, 3.84615384615385e01, 0, 0] @@ -30,7 +51,8 @@ ] -@pytest.mark.parametrize("material, dstrain, target", material_dstrain_stress_targets) -def test_compute_stress(material, dstrain, target): - stress = material.compute_stress(dstrain) +@pytest.mark.parametrize("particles, dstrain, target", particles_dstrain_stress_targets) +def test_compute_stress(particles, dstrain, target): + particles.dstrain = dstrain + stress = particles.material.compute_stress(particles) assert jnp.allclose(stress, target) From 763af134bf0d449cfee3ff6d7ca31fcd1645470a Mon Sep 17 00:00:00 2001 From: Chahak Mehta Date: Sat, 15 Jul 2023 16:59:12 -0700 Subject: [PATCH 04/21] Add Newtonian material --- diffmpm/materials/__init__.py | 1 + diffmpm/materials/newtonian.py | 115 ++++++++++++++++++++++++++++++++- tests/test_newtonian.py | 90 ++++++++++++++++++++++++++ 3 files changed, 205 insertions(+), 1 deletion(-) create mode 100644 tests/test_newtonian.py diff --git a/diffmpm/materials/__init__.py b/diffmpm/materials/__init__.py index ce35083..028e715 100644 --- a/diffmpm/materials/__init__.py +++ b/diffmpm/materials/__init__.py @@ -1,3 +1,4 @@ from diffmpm.materials._base import _Material from diffmpm.materials.simple import SimpleMaterial from diffmpm.materials.linear_elastic import LinearElastic +from diffmpm.materials.newtonian import Newtonian diff --git a/diffmpm/materials/newtonian.py b/diffmpm/materials/newtonian.py index e5a0d9b..558832f 100644 --- a/diffmpm/materials/newtonian.py +++ b/diffmpm/materials/newtonian.py @@ -1 +1,114 @@ -#!/usr/bin/env python3 +import jax.numpy as jnp +from jax import Array, lax +from jax.typing import ArrayLike + +from ._base import _Material + + +class Newtonian(_Material): + """Newtonian fluid material model.""" + + _props = ("density", "bulk_modulus", "dynamic_viscosity") + state_vars = ("pressure",) + + def __init__(self, material_properties: dict): + """Create a Newtonian material. + + Parameters + ---------- + material_properties: dict + Dictionary with material properties. For newtonian + materials, `density`, `bulk_modulus` and `dynamic_viscosity` + are required keys. + """ + self.validate_props(material_properties) + compressibility = 1 + + if material_properties.get("incompressible", False): + compressibility = 0 + + self.properties = { + **material_properties, + "compressibility": compressibility, + } + + def __repr__(self): + return f"Newtonian(props={self.properties})" + + def initialize_state_variables(self, nparticles: int) -> dict: + """Return initial state variables dictionary. + + Parameters + ---------- + nparticles : int + Number of particles being simulated with this material. + + Returns + ------- + dict + Dictionary of state variables initialized with values + decided by material type. + """ + state_vars_dict = {var: jnp.zeros((nparticles, 1)) for var in self.state_vars} + return state_vars_dict + + def _thermodynamic_pressure(self, volumetric_strain: ArrayLike) -> Array: + return -self.properties["bulk_modulus"] * volumetric_strain + + def compute_stress(self, particles): + """Compute material stress.""" + ndim = particles.loc.shape[-1] + if ndim not in {2, 3}: + raise ValueError(f"Cannot compute stress for {ndim}-d Newotonian material.") + volumetric_strain_rate = ( + particles.strain_rate[:, 0] + particles.strain_rate[:, 1] + ) + particles.state_vars["pressure"] = ( + particles.state_vars["pressure"] + .at[:] + .add( + self.properties["compressibility"] + * self._thermodynamic_pressure(particles.dvolumetric_strain) + ) + ) + + volumetric_stress_component = self.properties["compressibility"] * ( + -particles.state_vars["pressure"] + - (2 * self.properties["dynamic_viscosity"] * volumetric_strain_rate / 3) + ) + + stress = jnp.zeros_like(particles.stress) + stress = stress.at[:, 0].set( + volumetric_stress_component + + 2 * self.properties["dynamic_viscosity"] * particles.strain_rate[:, 0] + ) + stress = stress.at[:, 1].set( + volumetric_stress_component + + 2 * self.properties["dynamic_viscosity"] * particles.strain_rate[:, 1] + ) + + extra_component_2 = lax.select( + ndim == 3, + 2 * self.properties["dynamic_viscosity"] * particles.strain_rate[:, 2], + jnp.zeros_like(particles.strain_rate[:, 2]), + ) + stress = stress.at[:, 2].set(volumetric_stress_component + extra_component_2) + + stress = stress.at[:, 3].set( + self.properties["dynamic_viscosity"] * particles.strain_rate[:, 3] + ) + + component_4 = lax.select( + ndim == 3, + self.properties["dynamic_viscosity"] * particles.strain_rate[:, 4], + jnp.zeros_like(particles.strain_rate[:, 4]), + ) + stress = stress.at[:, 4].set(component_4) + component_5 = lax.select( + ndim == 3, + self.properties["dynamic_viscosity"] * particles.strain_rate[:, 5], + jnp.zeros_like(particles.strain_rate[:, 5]), + ) + stress = stress.at[:, 5].set(component_5) + + return stress diff --git a/tests/test_newtonian.py b/tests/test_newtonian.py new file mode 100644 index 0000000..518a246 --- /dev/null +++ b/tests/test_newtonian.py @@ -0,0 +1,90 @@ +import jax.numpy as jnp +import pytest +from diffmpm.constraint import Constraint +from diffmpm.element import Quadrilateral4Node +from diffmpm.materials import Newtonian +from diffmpm.node import Nodes +from diffmpm.particle import Particles + +particles_element_targets = [ + ( + Particles( + jnp.array([[0.5, 0.5]]).reshape(1, 1, 2), + Newtonian( + { + "density": 1000, + "bulk_modulus": 8333333.333333333, + "dynamic_viscosity": 8.9e-4, + } + ), + jnp.array([0]), + ), + Quadrilateral4Node( + (1, 1), + 1, + (4.0, 4.0), + [(0, Constraint(0, 0.02)), (0, Constraint(1, 0.03))], + Nodes(4, jnp.array([-2, -2, 2, -2, -2, 2, 2, 2]).reshape((4, 1, 2))), + ), + jnp.array( + [ + -52083.3333338896, + -52083.3333355583, + -52083.3333305521, + -0.0000041719, + 0, + 0, + ] + ).reshape(1, 6, 1), + ), + ( + Particles( + jnp.array([[0.5, 0.5]]).reshape(1, 1, 2), + Newtonian( + { + "density": 1000, + "bulk_modulus": 8333333.333333333, + "dynamic_viscosity": 8.9e-4, + "incompressible": True, + } + ), + jnp.array([0]), + ), + Quadrilateral4Node( + (1, 1), + 1, + (4.0, 4.0), + [(0, Constraint(0, 0.02)), (0, Constraint(1, 0.03))], + Nodes(4, jnp.array([-2, -2, 2, -2, -2, 2, 2, 2]).reshape((4, 1, 2))), + ), + jnp.array( + [ + -0.0000033375, + -0.00000500625, + 0, + -0.0000041719, + 0, + 0, + ] + ).reshape(1, 6, 1), + ), +] + + +@pytest.mark.parametrize( + "particles, element, target", + particles_element_targets, +) +def test_compute_stress(particles, element, target): + dt = 1 + particles.update_natural_coords(element) + if element.constraints: + element.apply_boundary_constraints() + particles.compute_strain(element, dt) + stress = particles.material.compute_stress(particles) + assert jnp.allclose(stress, target) + + +def test_init(): + with pytest.raises(KeyError): + Newtonian({"dynamic_viscosity": 1, "density": 1}) From 6120e44f38b1284141f8ecf5da2cd90bb8dabb57 Mon Sep 17 00:00:00 2001 From: Chahak Mehta Date: Tue, 18 Jul 2023 22:27:23 -0700 Subject: [PATCH 05/21] Convert list loops to tree_map calls and jit funcs --- diffmpm/element.py | 28 +++++++++++++++++++++++++--- diffmpm/mesh.py | 28 ++++++++++++++++++++++------ diffmpm/node.py | 2 ++ diffmpm/particle.py | 3 ++- 4 files changed, 51 insertions(+), 10 deletions(-) diff --git a/diffmpm/element.py b/diffmpm/element.py index 3eeff67..592afbc 100644 --- a/diffmpm/element.py +++ b/diffmpm/element.py @@ -2,6 +2,7 @@ import abc import itertools +from functools import partial from typing import TYPE_CHECKING, Optional, Sequence, Tuple if TYPE_CHECKING: @@ -9,10 +10,11 @@ import jax.numpy as jnp from jax import Array, jacobian, jit, lax, vmap -from jax.tree_util import register_pytree_node_class +from jax.tree_util import register_pytree_node_class, tree_map from jax.typing import ArrayLike from diffmpm.constraint import Constraint +from diffmpm.forces import NodalForce from diffmpm.node import Nodes __all__ = ["_Element", "Linear1D", "Quadrilateral4Node"] @@ -44,6 +46,7 @@ def id_to_node_ids(self, id: ArrayLike) -> Array: """ ... + @jit def id_to_node_loc(self, id: ArrayLike) -> Array: """Node locations corresponding to element `id`. @@ -61,6 +64,7 @@ def id_to_node_loc(self, id: ArrayLike) -> Array: node_ids = self.id_to_node_ids(id).squeeze() return self.nodes.loc[node_ids] + @jit def id_to_node_vel(self, id: ArrayLike) -> Array: """Node velocities corresponding to element `id`. @@ -135,6 +139,7 @@ def compute_nodal_mass(self, particles: Particles): Particles to map to the nodal values. """ + @jit def _step(pid, args): pmass, mass, mapped_pos, el_nodes = args mass = mass.at[el_nodes[pid]].add(pmass[pid] * mapped_pos[pid]) @@ -167,6 +172,7 @@ def compute_nodal_momentum(self, particles: Particles): Particles to map to the nodal values. """ + @jit def _step(pid, args): pmom, mom, mapped_pos, el_nodes = args mom = mom.at[el_nodes[pid]].add(mapped_pos[pid] @ pmom[pid]) @@ -217,6 +223,7 @@ def compute_external_force(self, particles: Particles): Particles to map to the nodal values. """ + @jit def _step(pid, args): f_ext, pf_ext, mapped_pos, el_nodes = args f_ext = f_ext.at[el_nodes[pid]].add(mapped_pos[pid] @ pf_ext[pid]) @@ -249,6 +256,7 @@ def compute_body_force(self, particles: Particles, gravity: ArrayLike): Particles to map to the nodal values. """ + @jit def _step(pid, args): f_ext, pmass, mapped_pos, el_nodes, gravity = args f_ext = f_ext.at[el_nodes[pid]].add( @@ -277,12 +285,20 @@ def apply_concentrated_nodal_forces(self, particles: Particles, curr_time: float curr_time: float Current time in the simulation. """ - for cnf in self.concentrated_nodal_forces: + + def _func(cnf, *, nodes): factor = cnf.function.value(curr_time) - self.nodes.f_ext = self.nodes.f_ext.at[cnf.node_ids, 0, cnf.dir].add( + nodes.f_ext = nodes.f_ext.at[cnf.node_ids, 0, cnf.dir].add( factor * cnf.force ) + partial_func = partial(_func, nodes=self.nodes) + tree_map( + partial_func, + self.concentrated_nodal_forces, + is_leaf=lambda x: isinstance(x, NodalForce), + ) + def apply_particle_traction_forces(self, particles: Particles): """Apply concentrated nodal forces. @@ -292,6 +308,7 @@ def apply_particle_traction_forces(self, particles: Particles): Particles in the simulation. """ + @jit def _step(pid, args): f_ext, ptraction, mapped_pos, el_nodes = args f_ext = f_ext.at[el_nodes[pid]].add(mapped_pos[pid] @ ptraction[pid]) @@ -689,6 +706,7 @@ def __init__( self.volume = jnp.asarray(volume) self.initialized = True + @jit def id_to_node_ids(self, id: ArrayLike): """Node IDs corresponding to element `id`. @@ -722,6 +740,7 @@ def id_to_node_ids(self, id: ArrayLike): ) return result.reshape(4, 1) + @jit def shapefn(self, xi: ArrayLike): """Evaluate linear shape function. @@ -756,6 +775,7 @@ def shapefn(self, xi: ArrayLike): result = result.transpose(1, 0, 2)[..., jnp.newaxis] return result + @jit def _shapefn_natural_grad(self, xi: ArrayLike): """Calculate the gradient of shape function. @@ -788,6 +808,7 @@ def _shapefn_natural_grad(self, xi: ArrayLike): ) return result + @jit def shapefn_grad(self, xi: ArrayLike, coords: ArrayLike): """Gradient of shape function in physical coordinates. @@ -857,6 +878,7 @@ def compute_internal_force(self, particles: Particles): Particles to map to the nodal values. """ + @jit def _step(pid, args): ( f_int, diff --git a/diffmpm/mesh.py b/diffmpm/mesh.py index 23bc6de..cf00fda 100644 --- a/diffmpm/mesh.py +++ b/diffmpm/mesh.py @@ -1,8 +1,10 @@ import abc +from functools import partial from typing import Callable, Sequence, Tuple import jax.numpy as jnp -from jax.tree_util import register_pytree_node_class +from jax import lax +from jax.tree_util import register_pytree_node_class, tree_map from diffmpm.element import _Element from diffmpm.particle import Particles @@ -39,8 +41,14 @@ def apply_on_elements(self, function: str, args: Tuple = ()): Parameters to be passed to the function. """ f = getattr(self.elements, function) - for particle_set in self.particles: - f(particle_set, *args) + + def _func(particles, *, func, fargs): + func(particles, *fargs) + + partial_func = partial(_func, func=f, fargs=args) + tree_map( + partial_func, self.particles, is_leaf=lambda x: isinstance(x, Particles) + ) # TODO: Convert to using jax directives for loop def apply_on_particles(self, function: str, args: Tuple = ()): @@ -53,9 +61,17 @@ def apply_on_particles(self, function: str, args: Tuple = ()): args: tuple Parameters to be passed to the function. """ - for particle_set in self.particles: - f = getattr(particle_set, function) - f(self.elements, *args) + + def _func(particles, *, elements, fname, fargs): + f = getattr(particles, fname) + f(elements, *fargs) + + partial_func = partial( + _func, elements=self.elements, fname=function, fargs=args + ) + tree_map( + partial_func, self.particles, is_leaf=lambda x: isinstance(x, Particles) + ) def apply_traction_on_particles(self, curr_time: float): """Apply tractions on particles. diff --git a/diffmpm/node.py b/diffmpm/node.py index 46e2a60..05d6a2e 100644 --- a/diffmpm/node.py +++ b/diffmpm/node.py @@ -1,6 +1,7 @@ from typing import Optional, Sized, Tuple import jax.numpy as jnp +from jax import jit from jax.tree_util import register_pytree_node_class from jax.typing import ArrayLike @@ -120,6 +121,7 @@ def __repr__(self): """Repr containing number of nodes.""" return f"Nodes(n={self.nnodes})" + @jit def get_total_force(self): """Calculate total force on the nodes.""" return self.f_int + self.f_ext + self.f_damp diff --git a/diffmpm/particle.py b/diffmpm/particle.py index 04f2581..bd7cdfa 100644 --- a/diffmpm/particle.py +++ b/diffmpm/particle.py @@ -1,7 +1,7 @@ from typing import Optional, Sized, Tuple import jax.numpy as jnp -from jax import lax, vmap +from jax import jit, lax, vmap from jax.tree_util import register_pytree_node_class from jax.typing import ArrayLike @@ -304,6 +304,7 @@ def _compute_strain_rate(self, dn_dx: ArrayLike, elements: _Element): temp = mapped_vel.squeeze(2) + @jit def _step(pid, args): dndx, nvel, strain_rate = args matmul = dndx[pid].T @ nvel[pid] From 0055007eb82bc5ef8195b34f0c8a2dcedbc1ebe0 Mon Sep 17 00:00:00 2001 From: Chahak Mehta Date: Fri, 21 Jul 2023 18:24:45 -0700 Subject: [PATCH 06/21] Run benchmark on cpu --- benchmarks/2d/uniaxial_particle_traction/test_benchmark.py | 6 ++++++ benchmarks/2d/uniaxial_stress/test_benchmark.py | 4 ++++ 2 files changed, 10 insertions(+) diff --git a/benchmarks/2d/uniaxial_particle_traction/test_benchmark.py b/benchmarks/2d/uniaxial_particle_traction/test_benchmark.py index 995ca16..1a53171 100644 --- a/benchmarks/2d/uniaxial_particle_traction/test_benchmark.py +++ b/benchmarks/2d/uniaxial_particle_traction/test_benchmark.py @@ -1,6 +1,9 @@ import os from pathlib import Path +import jax + +jax.config.update("jax_platform_name", "cpu") import jax.numpy as jnp from diffmpm.mpm import MPM @@ -31,3 +34,6 @@ def test_benchmarks(): result = jnp.load("results/uniaxial-particle-traction/particles_0990.npz") assert jnp.round(result["stress"][0, :, 0].min() - 0.750002924022295, 5) == 0.0 assert jnp.round(result["stress"][0, :, 0].max() - 0.9999997782938734, 5) == 0.0 + + +test_benchmarks() diff --git a/benchmarks/2d/uniaxial_stress/test_benchmark.py b/benchmarks/2d/uniaxial_stress/test_benchmark.py index 0dd6af8..ee589dd 100644 --- a/benchmarks/2d/uniaxial_stress/test_benchmark.py +++ b/benchmarks/2d/uniaxial_stress/test_benchmark.py @@ -1,6 +1,10 @@ import os from pathlib import Path +import jax + +jax.config.update("jax_platform_name", "cpu") + import jax.numpy as jnp from diffmpm.mpm import MPM From f28508819e9bb0640cdca3f434aebc1179fa187f Mon Sep 17 00:00:00 2001 From: Chahak Mehta Date: Fri, 21 Jul 2023 19:50:45 -0700 Subject: [PATCH 07/21] Temp commit: To be squashed --- .../uniaxial_nodal_forces/test_benchmark.py | 3 ++ .../test_benchmark.py | 3 -- diffmpm/element.py | 34 +++++++++---------- diffmpm/io.py | 16 +++++---- diffmpm/materials/__init__.py | 2 +- diffmpm/mesh.py | 24 +++++++++---- diffmpm/particle.py | 12 +++---- diffmpm/solver.py | 7 +++- diffmpm/writers.py | 4 +-- 9 files changed, 62 insertions(+), 43 deletions(-) diff --git a/benchmarks/2d/uniaxial_nodal_forces/test_benchmark.py b/benchmarks/2d/uniaxial_nodal_forces/test_benchmark.py index 4dcb077..5038150 100644 --- a/benchmarks/2d/uniaxial_nodal_forces/test_benchmark.py +++ b/benchmarks/2d/uniaxial_nodal_forces/test_benchmark.py @@ -1,6 +1,9 @@ import os from pathlib import Path +import jax + +jax.config.update("jax_platform_name", "cpu") import jax.numpy as jnp from diffmpm.mpm import MPM diff --git a/benchmarks/2d/uniaxial_particle_traction/test_benchmark.py b/benchmarks/2d/uniaxial_particle_traction/test_benchmark.py index 1a53171..5a46758 100644 --- a/benchmarks/2d/uniaxial_particle_traction/test_benchmark.py +++ b/benchmarks/2d/uniaxial_particle_traction/test_benchmark.py @@ -34,6 +34,3 @@ def test_benchmarks(): result = jnp.load("results/uniaxial-particle-traction/particles_0990.npz") assert jnp.round(result["stress"][0, :, 0].min() - 0.750002924022295, 5) == 0.0 assert jnp.round(result["stress"][0, :, 0].max() - 0.9999997782938734, 5) == 0.0 - - -test_benchmarks() diff --git a/diffmpm/element.py b/diffmpm/element.py index 592afbc..e252a7a 100644 --- a/diffmpm/element.py +++ b/diffmpm/element.py @@ -46,7 +46,7 @@ def id_to_node_ids(self, id: ArrayLike) -> Array: """ ... - @jit + # @jit def id_to_node_loc(self, id: ArrayLike) -> Array: """Node locations corresponding to element `id`. @@ -64,7 +64,7 @@ def id_to_node_loc(self, id: ArrayLike) -> Array: node_ids = self.id_to_node_ids(id).squeeze() return self.nodes.loc[node_ids] - @jit + # @jit def id_to_node_vel(self, id: ArrayLike) -> Array: """Node velocities corresponding to element `id`. @@ -147,7 +147,7 @@ def _step(pid, args): self.nodes.mass = self.nodes.mass.at[:].set(0) mapped_positions = self.shapefn(particles.reference_loc) - mapped_nodes = vmap(self.id_to_node_ids)(particles.element_ids).squeeze(-1) + mapped_nodes = vmap(jit(self.id_to_node_ids))(particles.element_ids).squeeze(-1) args = ( particles.mass, self.nodes.mass, @@ -180,7 +180,7 @@ def _step(pid, args): self.nodes.momentum = self.nodes.momentum.at[:].set(0) mapped_positions = self.shapefn(particles.reference_loc) - mapped_nodes = vmap(self.id_to_node_ids)(particles.element_ids).squeeze(-1) + mapped_nodes = vmap(jit(self.id_to_node_ids))(particles.element_ids).squeeze(-1) args = ( particles.mass * particles.velocity, self.nodes.momentum, @@ -231,7 +231,7 @@ def _step(pid, args): self.nodes.f_ext = self.nodes.f_ext.at[:].set(0) mapped_positions = self.shapefn(particles.reference_loc) - mapped_nodes = vmap(self.id_to_node_ids)(particles.element_ids).squeeze(-1) + mapped_nodes = vmap(jit(self.id_to_node_ids))(particles.element_ids).squeeze(-1) args = ( self.nodes.f_ext, particles.f_ext, @@ -265,7 +265,7 @@ def _step(pid, args): return f_ext, pmass, mapped_pos, el_nodes, gravity mapped_positions = self.shapefn(particles.reference_loc) - mapped_nodes = vmap(self.id_to_node_ids)(particles.element_ids).squeeze(-1) + mapped_nodes = vmap(jit(self.id_to_node_ids))(particles.element_ids).squeeze(-1) args = ( self.nodes.f_ext, particles.mass, @@ -315,7 +315,7 @@ def _step(pid, args): return f_ext, ptraction, mapped_pos, el_nodes mapped_positions = self.shapefn(particles.reference_loc) - mapped_nodes = vmap(self.id_to_node_ids)(particles.element_ids).squeeze(-1) + mapped_nodes = vmap(jit(self.id_to_node_ids))(particles.element_ids).squeeze(-1) args = (self.nodes.f_ext, particles.traction, mapped_positions, mapped_nodes) self.nodes.f_ext, _, _, _ = lax.fori_loop(0, len(particles), _step, args) @@ -601,9 +601,9 @@ def _step(pid, args): ) self.nodes.f_int = self.nodes.f_int.at[:].set(0) - mapped_nodes = vmap(self.id_to_node_ids)(particles.element_ids).squeeze(-1) - mapped_coords = vmap(self.id_to_node_loc)(particles.element_ids).squeeze(2) - mapped_grads = vmap(self.shapefn_grad)( + mapped_nodes = vmap(jit(self.id_to_node_ids))(particles.element_ids).squeeze(-1) + mapped_coords = vmap(jit(self.id_to_node_loc))(particles.element_ids).squeeze(2) + mapped_grads = vmap(jit(self.shapefn_grad))( particles.reference_loc[:, jnp.newaxis, ...], mapped_coords, ) @@ -706,7 +706,7 @@ def __init__( self.volume = jnp.asarray(volume) self.initialized = True - @jit + # @jit def id_to_node_ids(self, id: ArrayLike): """Node IDs corresponding to element `id`. @@ -740,7 +740,7 @@ def id_to_node_ids(self, id: ArrayLike): ) return result.reshape(4, 1) - @jit + # @jit def shapefn(self, xi: ArrayLike): """Evaluate linear shape function. @@ -775,7 +775,7 @@ def shapefn(self, xi: ArrayLike): result = result.transpose(1, 0, 2)[..., jnp.newaxis] return result - @jit + # @jit def _shapefn_natural_grad(self, xi: ArrayLike): """Calculate the gradient of shape function. @@ -808,7 +808,7 @@ def _shapefn_natural_grad(self, xi: ArrayLike): ) return result - @jit + # @jit def shapefn_grad(self, xi: ArrayLike, coords: ArrayLike): """Gradient of shape function in physical coordinates. @@ -907,9 +907,9 @@ def _step(pid, args): ) self.nodes.f_int = self.nodes.f_int.at[:].set(0) - mapped_nodes = vmap(self.id_to_node_ids)(particles.element_ids).squeeze(-1) - mapped_coords = vmap(self.id_to_node_loc)(particles.element_ids).squeeze(2) - mapped_grads = vmap(self.shapefn_grad)( + mapped_nodes = vmap(jit(self.id_to_node_ids))(particles.element_ids).squeeze(-1) + mapped_coords = vmap(jit(self.id_to_node_loc))(particles.element_ids).squeeze(2) + mapped_grads = vmap(jit(self.shapefn_grad))( particles.reference_loc[:, jnp.newaxis, ...], mapped_coords, ) diff --git a/diffmpm/io.py b/diffmpm/io.py index 66447f5..6cbc176 100644 --- a/diffmpm/io.py +++ b/diffmpm/io.py @@ -141,13 +141,15 @@ def _parse_mesh(self, config): element_cls = getattr(mpel, config["mesh"]["element"]) mesh_cls = getattr(mpmesh, f"Mesh{config['meta']['dimension']}D") - constraints = [ - ( - self._get_node_set_ids(c["nset_ids"]), - Constraint(c["dir"], c["velocity"]), - ) - for c in config["mesh"]["constraints"] - ] + constraints = [] + if "constraints" in config["mesh"]: + constraints = [ + ( + self._get_node_set_ids(c["nset_ids"]), + Constraint(c["dir"], c["velocity"]), + ) + for c in config["mesh"]["constraints"] + ] if config["mesh"]["type"] == "generator": elements = element_cls( diff --git a/diffmpm/materials/__init__.py b/diffmpm/materials/__init__.py index 028e715..8a023cd 100644 --- a/diffmpm/materials/__init__.py +++ b/diffmpm/materials/__init__.py @@ -1,4 +1,4 @@ from diffmpm.materials._base import _Material -from diffmpm.materials.simple import SimpleMaterial from diffmpm.materials.linear_elastic import LinearElastic from diffmpm.materials.newtonian import Newtonian +from diffmpm.materials.simple import SimpleMaterial diff --git a/diffmpm/mesh.py b/diffmpm/mesh.py index cf00fda..0176c5b 100644 --- a/diffmpm/mesh.py +++ b/diffmpm/mesh.py @@ -3,11 +3,12 @@ from typing import Callable, Sequence, Tuple import jax.numpy as jnp -from jax import lax +from jax import lax, jit from jax.tree_util import register_pytree_node_class, tree_map from diffmpm.element import _Element from diffmpm.particle import Particles +from diffmpm.forces import ParticleTraction __all__ = ["_MeshBase", "Mesh1D", "Mesh2D"] @@ -82,13 +83,24 @@ def apply_traction_on_particles(self, curr_time: float): Current time in the simulation. """ self.apply_on_particles("zero_traction") - for ptraction in self.particle_tractions: + + def func(ptraction, *, particle_sets): + def f(particles, *, ptraction, traction_val): + particles.assign_traction(ptraction.pids, ptraction.dir, traction_val) + factor = ptraction.function.value(curr_time) traction_val = factor * ptraction.traction - for i, pset_id in enumerate(ptraction.pset): - self.particles[pset_id].assign_traction( - ptraction.pids, ptraction.dir, traction_val - ) + partial_f = partial(f, ptraction=ptraction, traction_val=traction_val) + tree_map( + partial_f, particle_sets, is_leaf=lambda x: isinstance(x, Particles) + ) + + partial_func = partial(func, particle_sets=self.particles) + tree_map( + partial_func, + self.particle_tractions, + is_leaf=lambda x: isinstance(x, ParticleTraction), + ) self.apply_on_elements("apply_particle_traction_forces") diff --git a/diffmpm/particle.py b/diffmpm/particle.py index bd7cdfa..50dd991 100644 --- a/diffmpm/particle.py +++ b/diffmpm/particle.py @@ -206,7 +206,7 @@ def update_natural_coords(self, elements: _Element): Elements based on which to update the natural coordinates of the particles. """ - t = vmap(elements.id_to_node_loc)(self.element_ids) + t = vmap(jit(elements.id_to_node_loc))(self.element_ids) xi_coords = (self.loc - (t[:, 0, ...] + t[:, 2, ...]) / 2) * ( 2 / (t[:, 2, ...] - t[:, 0, ...]) ) @@ -231,7 +231,7 @@ def update_position_velocity( multiplied by dt. Default is False. """ mapped_positions = elements.shapefn(self.reference_loc) - mapped_ids = vmap(elements.id_to_node_ids)(self.element_ids).squeeze(-1) + mapped_ids = vmap(jit(elements.id_to_node_ids))(self.element_ids).squeeze(-1) nodal_velocity = jnp.sum( mapped_positions * elements.nodes.velocity[mapped_ids], axis=1 ) @@ -266,8 +266,8 @@ def compute_strain(self, elements: _Element, dt: float): dt : float Timestep. """ - mapped_coords = vmap(elements.id_to_node_loc)(self.element_ids).squeeze(2) - dn_dx_ = vmap(elements.shapefn_grad)( + mapped_coords = vmap(jit(elements.id_to_node_loc))(self.element_ids).squeeze(2) + dn_dx_ = vmap(jit(elements.shapefn_grad))( self.reference_loc[:, jnp.newaxis, ...], mapped_coords ) self.strain_rate = self._compute_strain_rate(dn_dx_, elements) @@ -275,7 +275,7 @@ def compute_strain(self, elements: _Element, dt: float): self.strain = self.strain.at[:].add(self.dstrain) centroids = jnp.zeros_like(self.loc) - dn_dx_centroid_ = vmap(elements.shapefn_grad)( + dn_dx_centroid_ = vmap(jit(elements.shapefn_grad))( centroids[:, jnp.newaxis, ...], mapped_coords ) strain_rate_centroid = self._compute_strain_rate(dn_dx_centroid_, elements) @@ -298,7 +298,7 @@ def _compute_strain_rate(self, dn_dx: ArrayLike, elements: _Element): """ dn_dx = jnp.asarray(dn_dx) strain_rate = jnp.zeros((dn_dx.shape[0], 6, 1)) # (nparticles, 6, 1) - mapped_vel = vmap(elements.id_to_node_vel)( + mapped_vel = vmap(jit(elements.id_to_node_vel))( self.element_ids ) # (nparticles, 2, 1) diff --git a/diffmpm/solver.py b/diffmpm/solver.py index 3b1ae01..8038d30 100644 --- a/diffmpm/solver.py +++ b/diffmpm/solver.py @@ -4,11 +4,12 @@ from typing import TYPE_CHECKING, Callable, Optional import jax.numpy as jnp -from jax import lax +from jax import lax, profiler from jax.experimental.host_callback import id_tap from jax.tree_util import register_pytree_node_class from jax.typing import ArrayLike +from diffmpm.pbar import loop_tqdm from diffmpm.scheme import USF, USL, _MPMScheme, _schemes if TYPE_CHECKING: @@ -142,6 +143,7 @@ def solve(self, gravity: ArrayLike): result = defaultdict(list) for step in tqdm(range(self.sim_steps)): + breakpoint() self.mpm_scheme.compute_nodal_kinematics() self.mpm_scheme.precompute_stress_strain() self.mpm_scheme.compute_forces(gravity, step) @@ -176,6 +178,7 @@ def solve_jit(self, gravity: ArrayLike) -> dict: final state of the simulation after completing all steps. """ + @loop_tqdm(self.sim_steps, print_rate=1) def _step(i, data): self = data self.mpm_scheme.compute_nodal_kinematics() @@ -210,6 +213,8 @@ def _write(self, i): ) return self + # with profiler.trace("/tmp/jax-trace", create_perfetto_link=True): + # self = lax.fori_loop(0, self.sim_steps, _step, self) self = lax.fori_loop(0, self.sim_steps, _step, self) arrays = {} for name in self.__particle_props: diff --git a/diffmpm/writers.py b/diffmpm/writers.py index fdc5cd2..b594fd7 100644 --- a/diffmpm/writers.py +++ b/diffmpm/writers.py @@ -1,10 +1,10 @@ import abc import logging from pathlib import Path +from typing import Annotated, Any, Tuple -from typing import Tuple, Annotated, Any -from jax.typing import ArrayLike import numpy as np +from jax.typing import ArrayLike logger = logging.getLogger(__file__) From 35183686a5c3504a79a1fdbb2cbc2455f9ac1e2f Mon Sep 17 00:00:00 2001 From: Chahak Mehta Date: Sat, 22 Jul 2023 14:18:16 -0700 Subject: [PATCH 08/21] Make Nodes as state --- diffmpm/constraint.py | 9 ++- diffmpm/element.py | 159 ++++++++++++++++++++++++------------------ diffmpm/node.py | 159 +++++++++++++----------------------------- tests/test_element.py | 15 ++-- 4 files changed, 156 insertions(+), 186 deletions(-) diff --git a/diffmpm/constraint.py b/diffmpm/constraint.py index 93f75bd..56dd038 100644 --- a/diffmpm/constraint.py +++ b/diffmpm/constraint.py @@ -37,8 +37,11 @@ def apply(self, obj, ids): The indices of the container `obj` on which the constraint will be applied. """ - obj.velocity = obj.velocity.at[ids, :, self.dir].set(self.velocity) - obj.momentum = obj.momentum.at[ids, :, self.dir].set( + velocity = obj.velocity.at[ids, :, self.dir].set(self.velocity) + momentum = obj.momentum.at[ids, :, self.dir].set( obj.mass[ids, :, 0] * self.velocity ) - obj.acceleration = obj.acceleration.at[ids, :, self.dir].set(0) + acceleration = obj.acceleration.at[ids, :, self.dir].set(0) + return obj.replace( + velocity=velocity, momentum=momentum, acceleration=acceleration + ) diff --git a/diffmpm/element.py b/diffmpm/element.py index e252a7a..adc9825 100644 --- a/diffmpm/element.py +++ b/diffmpm/element.py @@ -10,12 +10,12 @@ import jax.numpy as jnp from jax import Array, jacobian, jit, lax, vmap -from jax.tree_util import register_pytree_node_class, tree_map +from jax.tree_util import register_pytree_node_class, tree_map, tree_reduce from jax.typing import ArrayLike from diffmpm.constraint import Constraint from diffmpm.forces import NodalForce -from diffmpm.node import Nodes +from diffmpm.node import _NodesState, init_state __all__ = ["_Element", "Linear1D", "Quadrilateral4Node"] @@ -23,7 +23,7 @@ class _Element(abc.ABC): """Base element class that is inherited by all types of Elements.""" - nodes: Nodes + nodes: _NodesState total_elements: int concentrated_nodal_forces: Sequence volume: Array @@ -46,7 +46,6 @@ def id_to_node_ids(self, id: ArrayLike) -> Array: """ ... - # @jit def id_to_node_loc(self, id: ArrayLike) -> Array: """Node locations corresponding to element `id`. @@ -64,7 +63,6 @@ def id_to_node_loc(self, id: ArrayLike) -> Array: node_ids = self.id_to_node_ids(id).squeeze() return self.nodes.loc[node_ids] - # @jit def id_to_node_vel(self, id: ArrayLike) -> Array: """Node velocities corresponding to element `id`. @@ -145,16 +143,18 @@ def _step(pid, args): mass = mass.at[el_nodes[pid]].add(pmass[pid] * mapped_pos[pid]) return pmass, mass, mapped_pos, el_nodes - self.nodes.mass = self.nodes.mass.at[:].set(0) + mass = self.nodes.mass.at[:].set(0) mapped_positions = self.shapefn(particles.reference_loc) mapped_nodes = vmap(jit(self.id_to_node_ids))(particles.element_ids).squeeze(-1) args = ( particles.mass, - self.nodes.mass, + mass, mapped_positions, mapped_nodes, ) - _, self.nodes.mass, _, _ = lax.fori_loop(0, len(particles), _step, args) + _, mass, _, _ = lax.fori_loop(0, len(particles), _step, args) + # TODO: Return state instead of setting + self.nodes = self.nodes.replace(mass=mass) def compute_nodal_momentum(self, particles: Particles): r"""Compute the nodal mass based on particle mass. @@ -175,37 +175,37 @@ def compute_nodal_momentum(self, particles: Particles): @jit def _step(pid, args): pmom, mom, mapped_pos, el_nodes = args - mom = mom.at[el_nodes[pid]].add(mapped_pos[pid] @ pmom[pid]) - return pmom, mom, mapped_pos, el_nodes + new_mom = mom.at[el_nodes[pid]].add(mapped_pos[pid] @ pmom[pid]) + return pmom, new_mom, mapped_pos, el_nodes - self.nodes.momentum = self.nodes.momentum.at[:].set(0) + curr_mom = self.nodes.momentum.at[:].set(0) mapped_positions = self.shapefn(particles.reference_loc) mapped_nodes = vmap(jit(self.id_to_node_ids))(particles.element_ids).squeeze(-1) args = ( particles.mass * particles.velocity, - self.nodes.momentum, + curr_mom, mapped_positions, mapped_nodes, ) - _, self.nodes.momentum, _, _ = lax.fori_loop(0, len(particles), _step, args) - self.nodes.momentum = jnp.where( - jnp.abs(self.nodes.momentum) < 1e-12, - jnp.zeros_like(self.nodes.momentum), - self.nodes.momentum, - ) + _, new_momentum, _, _ = lax.fori_loop(0, len(particles), _step, args) + new_momentum = jnp.where(jnp.abs(new_momentum) < 1e-12, 0, new_momentum) + # TODO: Return state instead of setting + self.nodes = self.nodes.replace(momentum=new_momentum) def compute_velocity(self, particles: Particles): """Compute velocity using momentum.""" - self.nodes.velocity = jnp.where( + velocity = jnp.where( self.nodes.mass == 0, self.nodes.velocity, self.nodes.momentum / self.nodes.mass, ) - self.nodes.velocity = jnp.where( - jnp.abs(self.nodes.velocity) < 1e-12, - jnp.zeros_like(self.nodes.velocity), - self.nodes.velocity, + velocity = jnp.where( + jnp.abs(velocity) < 1e-12, + 0, + velocity, ) + # TODO: Return state instead of setting + self.nodes = self.nodes.replace(velocity=velocity) def compute_external_force(self, particles: Particles): r"""Update the nodal external force based on particle f_ext. @@ -229,16 +229,18 @@ def _step(pid, args): f_ext = f_ext.at[el_nodes[pid]].add(mapped_pos[pid] @ pf_ext[pid]) return f_ext, pf_ext, mapped_pos, el_nodes - self.nodes.f_ext = self.nodes.f_ext.at[:].set(0) + f_ext = self.nodes.f_ext.at[:].set(0) mapped_positions = self.shapefn(particles.reference_loc) mapped_nodes = vmap(jit(self.id_to_node_ids))(particles.element_ids).squeeze(-1) args = ( - self.nodes.f_ext, + f_ext, particles.f_ext, mapped_positions, mapped_nodes, ) - self.nodes.f_ext, _, _, _ = lax.fori_loop(0, len(particles), _step, args) + f_ext, _, _, _ = lax.fori_loop(0, len(particles), _step, args) + # TODO: Return state instead of setting + self.nodes = self.nodes.replace(f_ext=f_ext) def compute_body_force(self, particles: Particles, gravity: ArrayLike): r"""Update the nodal external force based on particle mass. @@ -273,7 +275,9 @@ def _step(pid, args): mapped_nodes, gravity, ) - self.nodes.f_ext, _, _, _, _ = lax.fori_loop(0, len(particles), _step, args) + f_ext, _, _, _, _ = lax.fori_loop(0, len(particles), _step, args) + # TODO: Return state instead of setting + self.nodes = self.nodes.replace(f_ext=f_ext) def apply_concentrated_nodal_forces(self, particles: Particles, curr_time: float): """Apply concentrated nodal forces. @@ -286,19 +290,22 @@ def apply_concentrated_nodal_forces(self, particles: Particles, curr_time: float Current time in the simulation. """ - def _func(cnf, *, nodes): + def _func(cnf, *, f_ext): factor = cnf.function.value(curr_time) - nodes.f_ext = nodes.f_ext.at[cnf.node_ids, 0, cnf.dir].add( - factor * cnf.force - ) + f_ext = f_ext.at[cnf.node_ids, 0, cnf.dir].add(factor * cnf.force) + return f_ext - partial_func = partial(_func, nodes=self.nodes) - tree_map( + partial_func = partial(_func, f_ext=self.nodes.f_ext) + _out = tree_map( partial_func, self.concentrated_nodal_forces, is_leaf=lambda x: isinstance(x, NodalForce), ) + f_ext = tree_reduce(lambda x, y: x + y, _out) + # TODO: Return state instead of setting + self.nodes = self.nodes.replace(f_ext=f_ext) + def apply_particle_traction_forces(self, particles: Particles): """Apply concentrated nodal forces. @@ -317,44 +324,64 @@ def _step(pid, args): mapped_positions = self.shapefn(particles.reference_loc) mapped_nodes = vmap(jit(self.id_to_node_ids))(particles.element_ids).squeeze(-1) args = (self.nodes.f_ext, particles.traction, mapped_positions, mapped_nodes) - self.nodes.f_ext, _, _, _ = lax.fori_loop(0, len(particles), _step, args) + f_ext, _, _, _ = lax.fori_loop(0, len(particles), _step, args) + # TODO: Return state instead of setting + self.nodes = self.nodes.replace(f_ext=f_ext) def update_nodal_acceleration_velocity( self, particles: Particles, dt: float, *args ): """Update the nodal momentum based on total force on nodes.""" - total_force = self.nodes.get_total_force() - self.nodes.acceleration = self.nodes.acceleration.at[:].set( + total_force = self.nodes.f_int + self.nodes.f_ext + self.nodes.f_damp + acceleration = self.nodes.acceleration.at[:].set( jnp.nan_to_num(jnp.divide(total_force, self.nodes.mass)) ) - self.nodes.velocity = self.nodes.velocity.at[:].add( - self.nodes.acceleration * dt - ) + velocity = self.nodes.velocity.at[:].add(acceleration * dt) self.apply_boundary_constraints() - self.nodes.momentum = self.nodes.momentum.at[:].set( - self.nodes.mass * self.nodes.velocity + momentum = self.nodes.momentum.at[:].set(self.nodes.mass * velocity) + velocity = jnp.where( + jnp.abs(velocity) < 1e-12, + 0, + velocity, ) - self.nodes.velocity = jnp.where( - jnp.abs(self.nodes.velocity) < 1e-12, - jnp.zeros_like(self.nodes.velocity), - self.nodes.velocity, + acceleration = jnp.where( + jnp.abs(acceleration) < 1e-12, + 0, + acceleration, ) - self.nodes.acceleration = jnp.where( - jnp.abs(self.nodes.acceleration) < 1e-12, - jnp.zeros_like(self.nodes.acceleration), - self.nodes.acceleration, + # TODO: Return state instead of setting + self.nodes = self.nodes.replace( + velocity=velocity, acceleration=acceleration, momentum=momentum ) def apply_boundary_constraints(self, *args): """Apply boundary conditions for nodal velocity.""" + + # TODO: Remove this for loop for ids, constraint in self.constraints: - constraint.apply(self.nodes, ids) + self.nodes = constraint.apply(self.nodes, ids) + + # def _func(i, args): + # constraints, state = args + # vel, mom, acc = constraints[i][1].apply(state, constraints[i][0]) + # new_state = state.replace(velocity=vel, momentum=mom, acceleration=acc) + # return constraints, new_state - def apply_force_boundary_constraints(self, *args): - """Apply boundary conditions for nodal forces.""" - self.nodes.f_int = self.nodes.f_int.at[self.constraints[0][0]].set(0) - self.nodes.f_ext = self.nodes.f_ext.at[self.constraints[0][0]].set(0) - self.nodes.f_damp = self.nodes.f_damp.at[self.constraints[0][0]].set(0) + # def _func2(constraint, *, nodes): + # return constraint[1].apply(nodes, constraint[0]) + + # from jax.tree_util import Partial + + # partial_func = Partial(_func2, nodes=self.nodes) + # _out = tree_map( + # partial_func, self.constraints, is_leaf=lambda x: isinstance(x, tuple) + # ) + + # breakpoint() + # _, self.nodes = lax.fori_loop( + # 0, len(self.constraints), _func, (self.constraints, self.nodes) + # ) + # breakpoint() @register_pytree_node_class @@ -646,7 +673,7 @@ def __init__( total_elements: int, el_len: float, constraints: Sequence[Tuple[ArrayLike, Constraint]], - nodes: Optional[Nodes] = None, + nodes: Optional[_NodesState] = None, concentrated_nodal_forces: Sequence = [], initialized: Optional[bool] = None, volume: Optional[ArrayLike] = None, @@ -694,7 +721,7 @@ def __init__( node_locations = ( jnp.asarray([coords[:, 1], coords[:, 0]]).T * self.el_len ).reshape(-1, 1, 2) - self.nodes = Nodes(int(total_nodes), node_locations) + self.nodes = init_state(int(total_nodes), node_locations) else: self.nodes = nodes @@ -706,7 +733,6 @@ def __init__( self.volume = jnp.asarray(volume) self.initialized = True - # @jit def id_to_node_ids(self, id: ArrayLike): """Node IDs corresponding to element `id`. @@ -740,7 +766,6 @@ def id_to_node_ids(self, id: ArrayLike): ) return result.reshape(4, 1) - # @jit def shapefn(self, xi: ArrayLike): """Evaluate linear shape function. @@ -775,7 +800,6 @@ def shapefn(self, xi: ArrayLike): result = result.transpose(1, 0, 2)[..., jnp.newaxis] return result - # @jit def _shapefn_natural_grad(self, xi: ArrayLike): """Calculate the gradient of shape function. @@ -808,7 +832,6 @@ def _shapefn_natural_grad(self, xi: ArrayLike): ) return result - # @jit def shapefn_grad(self, xi: ArrayLike, coords: ArrayLike): """Gradient of shape function in physical coordinates. @@ -906,21 +929,23 @@ def _step(pid, args): pstress, ) - self.nodes.f_int = self.nodes.f_int.at[:].set(0) - mapped_nodes = vmap(jit(self.id_to_node_ids))(particles.element_ids).squeeze(-1) - mapped_coords = vmap(jit(self.id_to_node_loc))(particles.element_ids).squeeze(2) - mapped_grads = vmap(jit(self.shapefn_grad))( + f_int = self.nodes.f_int.at[:].set(0) + mapped_nodes = vmap(self.id_to_node_ids)(particles.element_ids).squeeze(-1) + mapped_coords = vmap(self.id_to_node_loc)(particles.element_ids).squeeze(2) + mapped_grads = vmap(self.shapefn_grad)( particles.reference_loc[:, jnp.newaxis, ...], mapped_coords, ) args = ( - self.nodes.f_int, + f_int, particles.volume, mapped_grads, mapped_nodes, particles.stress, ) - self.nodes.f_int, _, _, _, _ = lax.fori_loop(0, len(particles), _step, args) + f_int, _, _, _, _ = lax.fori_loop(0, len(particles), _step, args) + # TODO: Return state instead of setting + self.nodes = self.nodes.replace(f_int=f_int) def compute_volume(self, *args): """Compute volume of all elements.""" diff --git a/diffmpm/node.py b/diffmpm/node.py index 05d6a2e..a035bd9 100644 --- a/diffmpm/node.py +++ b/diffmpm/node.py @@ -5,123 +5,62 @@ from jax.tree_util import register_pytree_node_class from jax.typing import ArrayLike +import chex -@register_pytree_node_class -class Nodes(Sized): - """Nodes container class. - Keeps track of all values required for nodal points. +@chex.dataclass(frozen=True) +class _NodesState: + nnodes: int + loc: chex.ArrayDevice + velocity: chex.ArrayDevice + acceleration: chex.ArrayDevice + mass: chex.ArrayDevice + momentum: chex.ArrayDevice + f_int: chex.ArrayDevice + f_ext: chex.ArrayDevice + f_damp: chex.ArrayDevice - Attributes + +def init_state( + nnodes: int, + loc: chex.ArrayDevice, +): + """Initialize container for Nodes. + + Parameters ---------- nnodes : int Number of nodes stored. loc : ArrayLike - Location of all the nodes. - velocity : array_like - Velocity of all the nodes. - mass : ArrayLike - Mass of all the nodes. - momentum : array_like - Momentum of all the nodes. - f_int : ArrayLike - Internal forces on all the nodes. - f_ext : ArrayLike - External forces present on all the nodes. - f_damp : ArrayLike - Damping forces on the nodes. + Locations of all the nodes. Expected shape (nnodes, 1, ndim) + initialized: bool + `False` if node property arrays like mass need to be initialized. + If `True`, they are set to values from `data`. + data: tuple + Tuple of length 7 that sets arrays for mass, density, volume, + and forces. Mainly used by JAX while unflattening. """ - - def __init__( - self, - nnodes: int, - loc: ArrayLike, - initialized: Optional[bool] = None, - data: Tuple[ArrayLike, ...] = tuple(), - ): - """Initialize container for Nodes. - - Parameters - ---------- - nnodes : int - Number of nodes stored. - loc : ArrayLike - Locations of all the nodes. Expected shape (nnodes, 1, ndim) - initialized: bool - `False` if node property arrays like mass need to be initialized. - If `True`, they are set to values from `data`. - data: tuple - Tuple of length 7 that sets arrays for mass, density, volume, - and forces. Mainly used by JAX while unflattening. - """ - self.nnodes = nnodes - loc = jnp.asarray(loc, dtype=jnp.float32) - if loc.ndim != 3: - raise ValueError( - f"`loc` should be of size (nnodes, 1, ndim); found {loc.shape}" - ) - self.loc = loc - - if initialized is None: - self.velocity = jnp.zeros_like(self.loc, dtype=jnp.float32) - self.acceleration = jnp.zeros_like(self.loc, dtype=jnp.float32) - self.mass = jnp.ones((self.loc.shape[0], 1, 1), dtype=jnp.float32) - self.momentum = jnp.zeros_like(self.loc, dtype=jnp.float32) - self.f_int = jnp.zeros_like(self.loc, dtype=jnp.float32) - self.f_ext = jnp.zeros_like(self.loc, dtype=jnp.float32) - self.f_damp = jnp.zeros_like(self.loc, dtype=jnp.float32) - else: - ( - self.velocity, - self.acceleration, - self.mass, - self.momentum, - self.f_int, - self.f_ext, - self.f_damp, - ) = data # type: ignore - self.initialized = True - - def tree_flatten(self): - """Flatten class as Pytree type.""" - children = ( - self.loc, - self.initialized, - self.velocity, - self.acceleration, - self.mass, - self.momentum, - self.f_int, - self.f_ext, - self.f_damp, + loc = jnp.asarray(loc, dtype=jnp.float32) + if loc.ndim != 3 or nnodes != loc.shape[0]: + raise ValueError( + f"`loc` should be of size (nnodes, 1, ndim); found {loc.shape}" ) - aux_data = (self.nnodes,) - return (children, aux_data) - - @classmethod - def tree_unflatten(cls, aux_data, children): - """Unflatten class from Pytree type.""" - return cls(aux_data[0], children[0], initialized=children[1], data=children[2:]) - - def reset_values(self): - """Reset nodal parameter values except location.""" - self.velocity = self.velocity.at[:].set(0) - self.acceleration = self.velocity.at[:].set(0) - self.mass = self.mass.at[:].set(0) - self.momentum = self.momentum.at[:].set(0) - self.f_int = self.f_int.at[:].set(0) - self.f_ext = self.f_ext.at[:].set(0) - self.f_damp = self.f_damp.at[:].set(0) - - def __len__(self): - """Set length of class as number of nodes.""" - return self.nnodes - - def __repr__(self): - """Repr containing number of nodes.""" - return f"Nodes(n={self.nnodes})" - @jit - def get_total_force(self): - """Calculate total force on the nodes.""" - return self.f_int + self.f_ext + self.f_damp + velocity = jnp.zeros_like(loc, dtype=jnp.float32) + acceleration = jnp.zeros_like(loc, dtype=jnp.float32) + mass = jnp.ones((loc.shape[0], 1, 1), dtype=jnp.float32) + momentum = jnp.zeros_like(loc, dtype=jnp.float32) + f_int = jnp.zeros_like(loc, dtype=jnp.float32) + f_ext = jnp.zeros_like(loc, dtype=jnp.float32) + f_damp = jnp.zeros_like(loc, dtype=jnp.float32) + return _NodesState( + nnodes=nnodes, + loc=loc, + velocity=velocity, + acceleration=acceleration, + mass=mass, + momentum=momentum, + f_int=f_int, + f_ext=f_ext, + f_damp=f_damp, + ) diff --git a/tests/test_element.py b/tests/test_element.py index ff8d92e..111609f 100644 --- a/tests/test_element.py +++ b/tests/test_element.py @@ -99,7 +99,9 @@ def test_element_node_loc(self, elements): assert jnp.all(node_loc == true_loc) def test_element_node_vel(self, elements): - elements.nodes.velocity += jnp.array([1, 1]) + elements.nodes = elements.nodes.replace( + velocity=elements.nodes.velocity + jnp.array([1, 1]) + ) node_vel = elements.id_to_node_vel(0) true_vel = jnp.array([[1.0, 1.0], [1, 1], [1, 1], [1, 1]]).reshape(4, 1, 2) assert jnp.all(node_vel == true_vel) @@ -160,18 +162,19 @@ def test_apply_concentrated_nodal_force(self, particles): ) def test_apply_boundary_constraints(self): - cons = [(jnp.array([0]), Constraint(0, 0))] + cons = [(jnp.array([0]), Constraint(0, 0)), (jnp.array([0]), Constraint(1, 2))] elements = Quadrilateral4Node((1, 1), 1, (1.0, 1.0), cons) - elements.nodes.velocity += 1 + elements.nodes = elements.nodes.replace(velocity=elements.nodes.velocity + 1) elements.apply_boundary_constraints() assert jnp.all( elements.nodes.velocity - == jnp.array([[0, 1], [1, 1], [1, 1], [1, 1]]).reshape(4, 1, 2) + == jnp.array([[0, 2], [1, 1], [1, 1], [1, 1]]).reshape(4, 1, 2) ) def test_update_nodal_acceleration_velocity(self, elements, particles): - elements.nodes.f_ext += jnp.array([1, 0]) - elements.nodes.mass = elements.nodes.mass.at[:].set(2) + f_ext = elements.nodes.f_ext + jnp.array([1, 0]) + mass = elements.nodes.mass.at[:].set(2) + elements.nodes = elements.nodes.replace(mass=mass, f_ext=f_ext) elements.update_nodal_acceleration_velocity(particles, 0.1) assert jnp.allclose( elements.nodes.acceleration, From 25289e7518e6eaa46f02756cf4e27c4bfc4ea8fd Mon Sep 17 00:00:00 2001 From: Chahak Mehta Date: Sat, 22 Jul 2023 16:01:34 -0700 Subject: [PATCH 09/21] Working nonfrozen node state --- diffmpm/element.py | 28 +++++++++++++++------------- diffmpm/node.py | 2 +- diffmpm/solver.py | 1 - 3 files changed, 16 insertions(+), 15 deletions(-) diff --git a/diffmpm/element.py b/diffmpm/element.py index adc9825..01b0cbb 100644 --- a/diffmpm/element.py +++ b/diffmpm/element.py @@ -295,16 +295,17 @@ def _func(cnf, *, f_ext): f_ext = f_ext.at[cnf.node_ids, 0, cnf.dir].add(factor * cnf.force) return f_ext - partial_func = partial(_func, f_ext=self.nodes.f_ext) - _out = tree_map( - partial_func, - self.concentrated_nodal_forces, - is_leaf=lambda x: isinstance(x, NodalForce), - ) + if self.concentrated_nodal_forces: + partial_func = partial(_func, f_ext=self.nodes.f_ext) + _out = tree_map( + partial_func, + self.concentrated_nodal_forces, + is_leaf=lambda x: isinstance(x, NodalForce), + ) - f_ext = tree_reduce(lambda x, y: x + y, _out) - # TODO: Return state instead of setting - self.nodes = self.nodes.replace(f_ext=f_ext) + f_ext = tree_reduce(lambda x, y: x + y, _out) + # TODO: Return state instead of setting + self.nodes = self.nodes.replace(f_ext=f_ext) def apply_particle_traction_forces(self, particles: Particles): """Apply concentrated nodal forces. @@ -337,17 +338,18 @@ def update_nodal_acceleration_velocity( jnp.nan_to_num(jnp.divide(total_force, self.nodes.mass)) ) velocity = self.nodes.velocity.at[:].add(acceleration * dt) + self.nodes = self.nodes.replace(velocity=velocity, acceleration=acceleration) self.apply_boundary_constraints() momentum = self.nodes.momentum.at[:].set(self.nodes.mass * velocity) velocity = jnp.where( - jnp.abs(velocity) < 1e-12, + jnp.abs(self.nodes.velocity) < 1e-12, 0, - velocity, + self.nodes.velocity, ) acceleration = jnp.where( - jnp.abs(acceleration) < 1e-12, + jnp.abs(self.nodes.acceleration) < 1e-12, 0, - acceleration, + self.nodes.acceleration, ) # TODO: Return state instead of setting self.nodes = self.nodes.replace( diff --git a/diffmpm/node.py b/diffmpm/node.py index a035bd9..a7c2490 100644 --- a/diffmpm/node.py +++ b/diffmpm/node.py @@ -8,7 +8,7 @@ import chex -@chex.dataclass(frozen=True) +@chex.dataclass() class _NodesState: nnodes: int loc: chex.ArrayDevice diff --git a/diffmpm/solver.py b/diffmpm/solver.py index 8038d30..7245dad 100644 --- a/diffmpm/solver.py +++ b/diffmpm/solver.py @@ -143,7 +143,6 @@ def solve(self, gravity: ArrayLike): result = defaultdict(list) for step in tqdm(range(self.sim_steps)): - breakpoint() self.mpm_scheme.compute_nodal_kinematics() self.mpm_scheme.precompute_stress_strain() self.mpm_scheme.compute_forces(gravity, step) From 00f9822e52f5f8c7d3d1b686a2d5f95cfbb89abb Mon Sep 17 00:00:00 2001 From: Chahak Mehta Date: Sat, 22 Jul 2023 17:35:52 -0700 Subject: [PATCH 10/21] Working frozen Nodestate --- diffmpm/constraint.py | 8 +++-- diffmpm/element.py | 68 +++++++++++++++++++++++++++++------------- diffmpm/node.py | 2 +- tests/test_element.py | 10 +++++-- tests/test_particle.py | 10 ++++--- 5 files changed, 66 insertions(+), 32 deletions(-) diff --git a/diffmpm/constraint.py b/diffmpm/constraint.py index 56dd038..2f8d7b5 100644 --- a/diffmpm/constraint.py +++ b/diffmpm/constraint.py @@ -42,6 +42,8 @@ def apply(self, obj, ids): obj.mass[ids, :, 0] * self.velocity ) acceleration = obj.acceleration.at[ids, :, self.dir].set(0) - return obj.replace( - velocity=velocity, momentum=momentum, acceleration=acceleration - ) + # return obj.replace( + # velocity=velocity, momentum=momentum, acceleration=acceleration + # ) + + return velocity, momentum, acceleration diff --git a/diffmpm/element.py b/diffmpm/element.py index 01b0cbb..3031867 100644 --- a/diffmpm/element.py +++ b/diffmpm/element.py @@ -11,6 +11,7 @@ import jax.numpy as jnp from jax import Array, jacobian, jit, lax, vmap from jax.tree_util import register_pytree_node_class, tree_map, tree_reduce +from jax import tree_util from jax.typing import ArrayLike from diffmpm.constraint import Constraint @@ -303,7 +304,14 @@ def _func(cnf, *, f_ext): is_leaf=lambda x: isinstance(x, NodalForce), ) - f_ext = tree_reduce(lambda x, y: x + y, _out) + def _f(x, *, orig): + return jnp.where(x == orig, 0, x) + + # This assumes that the nodal forces are not overlapping, i.e. + # no node will be acted by 2 forces in the same direction. + _step_1 = tree_map(partial(_f, orig=self.nodes.f_ext), _out) + _step_2 = tree_reduce(lambda x, y: x + y, _step_1) + f_ext = jnp.where(_step_2 == 0, self.nodes.f_ext, _step_2) # TODO: Return state instead of setting self.nodes = self.nodes.replace(f_ext=f_ext) @@ -359,31 +367,49 @@ def update_nodal_acceleration_velocity( def apply_boundary_constraints(self, *args): """Apply boundary conditions for nodal velocity.""" - # TODO: Remove this for loop - for ids, constraint in self.constraints: - self.nodes = constraint.apply(self.nodes, ids) + if self.constraints: - # def _func(i, args): - # constraints, state = args - # vel, mom, acc = constraints[i][1].apply(state, constraints[i][0]) - # new_state = state.replace(velocity=vel, momentum=mom, acceleration=acc) - # return constraints, new_state + def _func2(constraint, *, nodes): + return constraint[1].apply(nodes, constraint[0]) - # def _func2(constraint, *, nodes): - # return constraint[1].apply(nodes, constraint[0]) + partial_func = partial(_func2, nodes=self.nodes) + _out = tree_map( + partial_func, self.constraints, is_leaf=lambda x: isinstance(x, tuple) + ) - # from jax.tree_util import Partial + _temp = tree_util.tree_transpose( + tree_util.tree_structure([0 for e in _out]), + tree_util.tree_structure(_out[0]), + _out, + ) - # partial_func = Partial(_func2, nodes=self.nodes) - # _out = tree_map( - # partial_func, self.constraints, is_leaf=lambda x: isinstance(x, tuple) - # ) + def _f(x, *, orig): + return jnp.where(x == orig, jnp.nan, x) - # breakpoint() - # _, self.nodes = lax.fori_loop( - # 0, len(self.constraints), _func, (self.constraints, self.nodes) - # ) - # breakpoint() + _pf = partial(_f, orig=self.nodes.velocity) + _step_1 = tree_map(_pf, _temp[0]) + vel = tree_reduce( + lambda x, y: jnp.where(jnp.isnan(y), x, y), + [self.nodes.velocity, _step_1], + ) + + _pf = partial(_f, orig=self.nodes.momentum) + _step_1 = tree_map(_pf, _temp[1]) + mom = tree_reduce( + lambda x, y: jnp.where(jnp.isnan(y), x, y), + [self.nodes.momentum, _step_1], + ) + + _pf = partial(_f, orig=self.nodes.acceleration) + _step_1 = tree_map(_pf, _temp[2]) + acc = tree_reduce( + lambda x, y: jnp.where(jnp.isnan(y), x, y), + [self.nodes.acceleration, _step_1], + ) + # TODO: Return state instead of setting + self.nodes = self.nodes.replace( + velocity=vel, momentum=mom, acceleration=acc + ) @register_pytree_node_class diff --git a/diffmpm/node.py b/diffmpm/node.py index a7c2490..a035bd9 100644 --- a/diffmpm/node.py +++ b/diffmpm/node.py @@ -8,7 +8,7 @@ import chex -@chex.dataclass() +@chex.dataclass(frozen=True) class _NodesState: nnodes: int loc: chex.ArrayDevice diff --git a/tests/test_element.py b/tests/test_element.py index 111609f..11e9d6a 100644 --- a/tests/test_element.py +++ b/tests/test_element.py @@ -155,20 +155,24 @@ def test_apply_concentrated_nodal_force(self, particles): elements = Quadrilateral4Node( (1, 1), 1, 1, [], concentrated_nodal_forces=[cnf_1, cnf_2] ) + elements.nodes = elements.nodes.replace(f_ext=elements.nodes.f_ext + 2) elements.apply_concentrated_nodal_forces(particles, 1) assert jnp.all( elements.nodes.f_ext - == jnp.array([[1, 0], [0, 0], [1, 1], [0, 0]]).reshape(4, 1, 2) + == jnp.array([[3, 2], [2, 2], [3, 3], [2, 2]]).reshape(4, 1, 2) ) def test_apply_boundary_constraints(self): - cons = [(jnp.array([0]), Constraint(0, 0)), (jnp.array([0]), Constraint(1, 2))] + cons = [ + (jnp.array([0, 1]), Constraint(0, 0)), + (jnp.array([0]), Constraint(1, 2)), + ] elements = Quadrilateral4Node((1, 1), 1, (1.0, 1.0), cons) elements.nodes = elements.nodes.replace(velocity=elements.nodes.velocity + 1) elements.apply_boundary_constraints() assert jnp.all( elements.nodes.velocity - == jnp.array([[0, 2], [1, 1], [1, 1], [1, 1]]).reshape(4, 1, 2) + == jnp.array([[0, 2], [0, 1], [1, 1], [1, 1]]).reshape(4, 1, 2) ) def test_update_nodal_acceleration_velocity(self, elements, particles): diff --git a/tests/test_particle.py b/tests/test_particle.py index d67bc2f..07b8c19 100644 --- a/tests/test_particle.py +++ b/tests/test_particle.py @@ -26,14 +26,16 @@ def particles(self): ) def test_update_velocity(self, elements, particles, velocity_update, expected): particles.update_natural_coords(elements) - elements.nodes.acceleration += 1 - elements.nodes.velocity += 1 + elements.nodes = elements.nodes.replace( + acceleration=elements.nodes.acceleration + 1 + ) + elements.nodes = elements.nodes.replace(velocity=elements.nodes.velocity + 1) particles.update_position_velocity(elements, 0.1, velocity_update) assert jnp.allclose(particles.velocity, expected) def test_compute_strain(self, elements, particles): - elements.nodes.velocity = jnp.array([[0, 1], [0, 2], [0, 3], [0, 4]]).reshape( - 4, 1, 2 + elements.nodes = elements.nodes.replace( + velocity=jnp.array([[0, 1], [0, 2], [0, 3], [0, 4]]).reshape(4, 1, 2) ) particles.update_natural_coords(elements) particles.compute_strain(elements, 0.1) From 28e7bff5f4ff699c348f279a5980059adc34825b Mon Sep 17 00:00:00 2001 From: Chahak Mehta Date: Sat, 22 Jul 2023 21:10:16 -0700 Subject: [PATCH 11/21] Working nonfrozen particle state --- .../mpm-nodal-forces.toml | 2 +- .../mpm-particle-traction.toml | 2 +- .../test_benchmark.py | 3 + .../uniaxial_stress/mpm-uniaxial-stress.toml | 2 +- diffmpm/element.py | 16 +- diffmpm/io.py | 11 +- diffmpm/materials/__init__.py | 8 +- diffmpm/materials/linear_elastic.py | 187 +++++++--- diffmpm/materials/simple.py | 17 + diffmpm/mesh.py | 60 ++- diffmpm/particle.py | 345 ++++++++++++++++++ tests/{test_newtonian.py => newtonian.py} | 0 tests/test_element.py | 22 +- tests/test_material.py | 26 +- tests/test_particle.py | 31 +- 15 files changed, 601 insertions(+), 131 deletions(-) rename tests/{test_newtonian.py => newtonian.py} (100%) diff --git a/benchmarks/2d/uniaxial_nodal_forces/mpm-nodal-forces.toml b/benchmarks/2d/uniaxial_nodal_forces/mpm-nodal-forces.toml index 1e7ef1a..93a9386 100644 --- a/benchmarks/2d/uniaxial_nodal_forces/mpm-nodal-forces.toml +++ b/benchmarks/2d/uniaxial_nodal_forces/mpm-nodal-forces.toml @@ -46,7 +46,7 @@ id = 0 density = 1000 poisson_ratio = 0 youngs_modulus = 1000000 -type = "LinearElastic" +type = "linear_elastic" [[particles]] file = "particles-2d-nodal-force.json" diff --git a/benchmarks/2d/uniaxial_particle_traction/mpm-particle-traction.toml b/benchmarks/2d/uniaxial_particle_traction/mpm-particle-traction.toml index 480ec4e..067cbb9 100644 --- a/benchmarks/2d/uniaxial_particle_traction/mpm-particle-traction.toml +++ b/benchmarks/2d/uniaxial_particle_traction/mpm-particle-traction.toml @@ -46,7 +46,7 @@ id = 0 density = 1000 poisson_ratio = 0 youngs_modulus = 1000000 -type = "LinearElastic" +type = "linear_elastic" [[particles]] file = "particles-2d-particle-traction.json" diff --git a/benchmarks/2d/uniaxial_particle_traction/test_benchmark.py b/benchmarks/2d/uniaxial_particle_traction/test_benchmark.py index 5a46758..1a53171 100644 --- a/benchmarks/2d/uniaxial_particle_traction/test_benchmark.py +++ b/benchmarks/2d/uniaxial_particle_traction/test_benchmark.py @@ -34,3 +34,6 @@ def test_benchmarks(): result = jnp.load("results/uniaxial-particle-traction/particles_0990.npz") assert jnp.round(result["stress"][0, :, 0].min() - 0.750002924022295, 5) == 0.0 assert jnp.round(result["stress"][0, :, 0].max() - 0.9999997782938734, 5) == 0.0 + + +test_benchmarks() diff --git a/benchmarks/2d/uniaxial_stress/mpm-uniaxial-stress.toml b/benchmarks/2d/uniaxial_stress/mpm-uniaxial-stress.toml index 3e074cd..bf15148 100644 --- a/benchmarks/2d/uniaxial_stress/mpm-uniaxial-stress.toml +++ b/benchmarks/2d/uniaxial_stress/mpm-uniaxial-stress.toml @@ -47,7 +47,7 @@ id = 0 density = 1 poisson_ratio = 0 youngs_modulus = 1000 -type = "LinearElastic" +type = "linear_elastic" [[particles]] file = "particles-2d-uniaxial-stress.json" diff --git a/diffmpm/element.py b/diffmpm/element.py index 3031867..bf3d0d6 100644 --- a/diffmpm/element.py +++ b/diffmpm/element.py @@ -153,7 +153,7 @@ def _step(pid, args): mapped_positions, mapped_nodes, ) - _, mass, _, _ = lax.fori_loop(0, len(particles), _step, args) + _, mass, _, _ = lax.fori_loop(0, particles.nparticles, _step, args) # TODO: Return state instead of setting self.nodes = self.nodes.replace(mass=mass) @@ -188,7 +188,7 @@ def _step(pid, args): mapped_positions, mapped_nodes, ) - _, new_momentum, _, _ = lax.fori_loop(0, len(particles), _step, args) + _, new_momentum, _, _ = lax.fori_loop(0, particles.nparticles, _step, args) new_momentum = jnp.where(jnp.abs(new_momentum) < 1e-12, 0, new_momentum) # TODO: Return state instead of setting self.nodes = self.nodes.replace(momentum=new_momentum) @@ -239,7 +239,7 @@ def _step(pid, args): mapped_positions, mapped_nodes, ) - f_ext, _, _, _ = lax.fori_loop(0, len(particles), _step, args) + f_ext, _, _, _ = lax.fori_loop(0, particles.nparticles, _step, args) # TODO: Return state instead of setting self.nodes = self.nodes.replace(f_ext=f_ext) @@ -276,7 +276,7 @@ def _step(pid, args): mapped_nodes, gravity, ) - f_ext, _, _, _, _ = lax.fori_loop(0, len(particles), _step, args) + f_ext, _, _, _, _ = lax.fori_loop(0, particles.nparticles, _step, args) # TODO: Return state instead of setting self.nodes = self.nodes.replace(f_ext=f_ext) @@ -333,7 +333,7 @@ def _step(pid, args): mapped_positions = self.shapefn(particles.reference_loc) mapped_nodes = vmap(jit(self.id_to_node_ids))(particles.element_ids).squeeze(-1) args = (self.nodes.f_ext, particles.traction, mapped_positions, mapped_nodes) - f_ext, _, _, _ = lax.fori_loop(0, len(particles), _step, args) + f_ext, _, _, _ = lax.fori_loop(0, particles.nparticles, _step, args) # TODO: Return state instead of setting self.nodes = self.nodes.replace(f_ext=f_ext) @@ -669,7 +669,9 @@ def _step(pid, args): mapped_nodes, particles.stress, ) - self.nodes.f_int, _, _, _, _ = lax.fori_loop(0, len(particles), _step, args) + self.nodes.f_int, _, _, _, _ = lax.fori_loop( + 0, particles.nparticles, _step, args + ) @register_pytree_node_class @@ -971,7 +973,7 @@ def _step(pid, args): mapped_nodes, particles.stress, ) - f_int, _, _, _, _ = lax.fori_loop(0, len(particles), _step, args) + f_int, _, _, _, _ = lax.fori_loop(0, particles.nparticles, _step, args) # TODO: Return state instead of setting self.nodes = self.nodes.replace(f_int=f_int) diff --git a/diffmpm/io.py b/diffmpm/io.py index 6cbc176..6a4f46e 100644 --- a/diffmpm/io.py +++ b/diffmpm/io.py @@ -9,7 +9,7 @@ from diffmpm.constraint import Constraint from diffmpm.forces import NodalForce, ParticleTraction from diffmpm.functions import Linear, Unit -from diffmpm.particle import Particles +from diffmpm.particle import Particles, _ParticlesState, init_particle_state class Config: @@ -58,9 +58,10 @@ def _parse_materials(self, config): materials = [] for mat_config in config["materials"]: mat_type = mat_config.pop("type") - mat_cls = getattr(mpmat, mat_type) - mat = mat_cls(mat_config) - materials.append(mat) + # mat_cls = getattr(mpmat, mat_type) + # mat = mat_cls(mat_config) + mat_fun = getattr(mpmat, f"init_{mat_type}") + materials.append(mat_fun(mat_config)) self.parsed_config["materials"] = materials def _parse_particles(self, config): @@ -70,7 +71,7 @@ def _parse_particles(self, config): with open(pset_config["file"], "r") as f: ploc = jnp.asarray(json.load(f)) peids = jnp.zeros(ploc.shape[0], dtype=jnp.int32) - pset = Particles(ploc, pmat, peids) + pset = init_particle_state(ploc, pmat, peids) pset.velocity = pset.velocity.at[:].set(pset_config["init_velocity"]) particle_sets.append(pset) self.parsed_config["particles"] = particle_sets diff --git a/diffmpm/materials/__init__.py b/diffmpm/materials/__init__.py index 8a023cd..a7bc501 100644 --- a/diffmpm/materials/__init__.py +++ b/diffmpm/materials/__init__.py @@ -1,4 +1,8 @@ from diffmpm.materials._base import _Material -from diffmpm.materials.linear_elastic import LinearElastic + +# from diffmpm.materials.linear_elastic import LinearElastic +from diffmpm.materials.linear_elastic import _LinearElasticState, init_linear_elastic from diffmpm.materials.newtonian import Newtonian -from diffmpm.materials.simple import SimpleMaterial + +# from diffmpm.materials.simple import SimpleMaterial +from diffmpm.materials.simple import _SimpleMaterialState, init_simple diff --git a/diffmpm/materials/linear_elastic.py b/diffmpm/materials/linear_elastic.py index 5a008d4..7e9c543 100644 --- a/diffmpm/materials/linear_elastic.py +++ b/diffmpm/materials/linear_elastic.py @@ -3,68 +3,135 @@ from ._base import _Material +import chex -@register_pytree_node_class -class LinearElastic(_Material): - """Linear Elastic Material.""" - _props = ("density", "youngs_modulus", "poisson_ratio") - state_vars = () +@chex.dataclass() +class _LinearElasticState: + id: int + state_vars: tuple + density: float + youngs_modulus: float + poisson_ratio: float + bulk_modulus: float + pwave_velocity: float + swave_velocity: float + de: chex.ArrayDevice - def __init__(self, material_properties): - """Create a Linear Elastic material. - - Parameters - ---------- - material_properties: dict - Dictionary with material properties. For linear elastic - materials, 'density' and 'youngs_modulus' are required keys. - """ - self.validate_props(material_properties) - youngs_modulus = material_properties["youngs_modulus"] - poisson_ratio = material_properties["poisson_ratio"] - density = material_properties["density"] - bulk_modulus = youngs_modulus / (3 * (1 - 2 * poisson_ratio)) - constrained_modulus = ( - youngs_modulus - * (1 - poisson_ratio) - / ((1 + poisson_ratio) * (1 - 2 * poisson_ratio)) - ) - shear_modulus = youngs_modulus / (2 * (1 + poisson_ratio)) - # Wave velocities - vp = jnp.sqrt(constrained_modulus / density) - vs = jnp.sqrt(shear_modulus / density) - self.properties = { - **material_properties, - "bulk_modulus": bulk_modulus, - "pwave_velocity": vp, - "swave_velocity": vs, - } - self._compute_elastic_tensor() - - def __repr__(self): - return f"LinearElastic(props={self.properties})" - - def _compute_elastic_tensor(self): - G = self.properties["youngs_modulus"] / ( - 2 * (1 + self.properties["poisson_ratio"]) - ) - - a1 = self.properties["bulk_modulus"] + (4 * G / 3) - a2 = self.properties["bulk_modulus"] - (2 * G / 3) - - self.de = jnp.array( - [ - [a1, a2, a2, 0, 0, 0], - [a2, a1, a2, 0, 0, 0], - [a2, a2, a1, 0, 0, 0], - [0, 0, 0, G, 0, 0], - [0, 0, 0, 0, G, 0], - [0, 0, 0, 0, 0, G], - ] - ) - - def compute_stress(self, particles): + def compute_stress(self, state): """Compute material stress.""" - dstress = self.de @ particles.dstrain + dstress = self.de @ state.dstrain return dstress + + +def init_linear_elastic(material_properties): + """Create a Linear Elastic material. + + Parameters + ---------- + material_properties: dict + Dictionary with material properties. For linear elastic + materials, 'density' and 'youngs_modulus' are required keys. + """ + state_vars = () + youngs_modulus = material_properties["youngs_modulus"] + poisson_ratio = material_properties["poisson_ratio"] + density = material_properties["density"] + bulk_modulus = youngs_modulus / (3 * (1 - 2 * poisson_ratio)) + constrained_modulus = ( + youngs_modulus + * (1 - poisson_ratio) + / ((1 + poisson_ratio) * (1 - 2 * poisson_ratio)) + ) + shear_modulus = youngs_modulus / (2 * (1 + poisson_ratio)) + # Wave velocities + vp = jnp.sqrt(constrained_modulus / density) + vs = jnp.sqrt(shear_modulus / density) + properties = { + **material_properties, + "bulk_modulus": bulk_modulus, + "pwave_velocity": vp, + "swave_velocity": vs, + } + G = youngs_modulus / (2 * (1 + poisson_ratio)) + + a1 = bulk_modulus + (4 * G / 3) + a2 = bulk_modulus - (2 * G / 3) + + de = jnp.array( + [ + [a1, a2, a2, 0, 0, 0], + [a2, a1, a2, 0, 0, 0], + [a2, a2, a1, 0, 0, 0], + [0, 0, 0, G, 0, 0], + [0, 0, 0, 0, G, 0], + [0, 0, 0, 0, 0, G], + ] + ) + return _LinearElasticState(**properties, de=de, state_vars=state_vars) + + +# @register_pytree_node_class +# class LinearElastic(_Material): +# """Linear Elastic Material.""" + +# _props = ("density", "youngs_modulus", "poisson_ratio") +# state_vars = () + +# def __init__(self, material_properties): +# """Create a Linear Elastic material. + +# Parameters +# ---------- +# material_properties: dict +# Dictionary with material properties. For linear elastic +# materials, 'density' and 'youngs_modulus' are required keys. +# """ +# self.validate_props(material_properties) +# youngs_modulus = material_properties["youngs_modulus"] +# poisson_ratio = material_properties["poisson_ratio"] +# density = material_properties["density"] +# bulk_modulus = youngs_modulus / (3 * (1 - 2 * poisson_ratio)) +# constrained_modulus = ( +# youngs_modulus +# * (1 - poisson_ratio) +# / ((1 + poisson_ratio) * (1 - 2 * poisson_ratio)) +# ) +# shear_modulus = youngs_modulus / (2 * (1 + poisson_ratio)) +# # Wave velocities +# vp = jnp.sqrt(constrained_modulus / density) +# vs = jnp.sqrt(shear_modulus / density) +# self.properties = { +# **material_properties, +# "bulk_modulus": bulk_modulus, +# "pwave_velocity": vp, +# "swave_velocity": vs, +# } +# self._compute_elastic_tensor() + +# def __repr__(self): +# return f"LinearElastic(props={self.properties})" + +# def _compute_elastic_tensor(self): +# G = self.properties["youngs_modulus"] / ( +# 2 * (1 + self.properties["poisson_ratio"]) +# ) + +# a1 = self.properties["bulk_modulus"] + (4 * G / 3) +# a2 = self.properties["bulk_modulus"] - (2 * G / 3) + +# self.de = jnp.array( +# [ +# [a1, a2, a2, 0, 0, 0], +# [a2, a1, a2, 0, 0, 0], +# [a2, a2, a1, 0, 0, 0], +# [0, 0, 0, G, 0, 0], +# [0, 0, 0, 0, G, 0], +# [0, 0, 0, 0, 0, G], +# ] +# ) + +# def compute_stress(self, particles): +# """Compute material stress.""" +# dstress = self.de @ particles.dstrain +# return dstress diff --git a/diffmpm/materials/simple.py b/diffmpm/materials/simple.py index 77b57ca..349190d 100644 --- a/diffmpm/materials/simple.py +++ b/diffmpm/materials/simple.py @@ -2,6 +2,23 @@ from ._base import _Material +import chex + + +@chex.dataclass() +class _SimpleMaterialState: + id: int + E: float + density: float + state_vars: () + + def compute_stress(self, particles): + return particles.dstrain * self.E + + +def init_simple(material_properties): + return _SimpleMaterialState(**material_properties, state_vars=()) + @register_pytree_node_class class SimpleMaterial(_Material): diff --git a/diffmpm/mesh.py b/diffmpm/mesh.py index 0176c5b..ce68d6e 100644 --- a/diffmpm/mesh.py +++ b/diffmpm/mesh.py @@ -3,11 +3,12 @@ from typing import Callable, Sequence, Tuple import jax.numpy as jnp -from jax import lax, jit +from jax import lax, jit, tree_util from jax.tree_util import register_pytree_node_class, tree_map from diffmpm.element import _Element -from diffmpm.particle import Particles +from diffmpm.particle import Particles, _ParticlesState +import diffmpm.particle as dpart from diffmpm.forces import ParticleTraction __all__ = ["_MeshBase", "Mesh1D", "Mesh2D"] @@ -26,7 +27,7 @@ class _MeshBase(abc.ABC): def __init__(self, config: dict): """Initialize mesh using configuration.""" - self.particles: Sequence[Particles] = config["particles"] + self.particles: Sequence[_ParticlesState] = config["particles"] self.elements: _Element = config["elements"] self.particle_tractions = config["particle_surface_traction"] @@ -48,7 +49,9 @@ def _func(particles, *, func, fargs): partial_func = partial(_func, func=f, fargs=args) tree_map( - partial_func, self.particles, is_leaf=lambda x: isinstance(x, Particles) + partial_func, + self.particles, + is_leaf=lambda x: isinstance(x, _ParticlesState), ) # TODO: Convert to using jax directives for loop @@ -58,21 +61,24 @@ def apply_on_particles(self, function: str, args: Tuple = ()): Parameters ---------- function: str - A string corresponding to a function name in `Particles`. + A string corresponding to a function name in `_ParticlesState`. args: tuple Parameters to be passed to the function. """ def _func(particles, *, elements, fname, fargs): - f = getattr(particles, fname) - f(elements, *fargs) + f = getattr(dpart, fname) + return f(particles, elements, *fargs) partial_func = partial( _func, elements=self.elements, fname=function, fargs=args ) - tree_map( - partial_func, self.particles, is_leaf=lambda x: isinstance(x, Particles) + new_states = tree_map( + partial_func, + self.particles, + is_leaf=lambda x: isinstance(x, _ParticlesState), ) + self.particles = new_states def apply_traction_on_particles(self, curr_time: float): """Apply tractions on particles. @@ -86,23 +92,41 @@ def apply_traction_on_particles(self, curr_time: float): def func(ptraction, *, particle_sets): def f(particles, *, ptraction, traction_val): - particles.assign_traction(ptraction.pids, ptraction.dir, traction_val) + return dpart.assign_traction( + particles, ptraction.pids, ptraction.dir, traction_val + ) factor = ptraction.function.value(curr_time) traction_val = factor * ptraction.traction partial_f = partial(f, ptraction=ptraction, traction_val=traction_val) - tree_map( - partial_f, particle_sets, is_leaf=lambda x: isinstance(x, Particles) + traction_sets = tree_map( + partial_f, + particle_sets, + is_leaf=lambda x: isinstance(x, _ParticlesState), ) + return tuple(traction_sets) partial_func = partial(func, particle_sets=self.particles) - tree_map( - partial_func, - self.particle_tractions, - is_leaf=lambda x: isinstance(x, ParticleTraction), - ) + if self.particle_tractions: + _out = tree_map( + partial_func, + self.particle_tractions, + is_leaf=lambda x: isinstance(x, ParticleTraction), + ) + _temp = tree_util.tree_transpose( + tree_util.tree_structure([0 for e in _out]), + tree_util.tree_structure(_out[0]), + _out, + ) + tractions_ = tree_util.tree_reduce( + lambda x, y: x + y, _temp, is_leaf=lambda x: isinstance(x, list) + ) + self.particles = [ + pset.replace(traction=traction) + for pset, traction in zip(self.particles, tractions_) + ] - self.apply_on_elements("apply_particle_traction_forces") + self.apply_on_elements("apply_particle_traction_forces") def tree_flatten(self): children = (self.particles, self.elements) diff --git a/diffmpm/particle.py b/diffmpm/particle.py index 50dd991..bbe17a2 100644 --- a/diffmpm/particle.py +++ b/diffmpm/particle.py @@ -8,6 +8,351 @@ from diffmpm.element import _Element from diffmpm.materials import _Material +import chex + + +@chex.dataclass() +class _ParticlesState: + nparticles: int + loc: chex.ArrayDevice + material: _Material + element_ids: chex.ArrayDevice + mass: chex.ArrayDevice + density: chex.ArrayDevice + volume: chex.ArrayDevice + size: chex.ArrayDevice + velocity: chex.ArrayDevice + acceleration: chex.ArrayDevice + momentum: chex.ArrayDevice + strain: chex.ArrayDevice + stress: chex.ArrayDevice + strain_rate: chex.ArrayDevice + dstrain: chex.ArrayDevice + f_ext: chex.ArrayDevice + traction: chex.ArrayDevice + reference_loc: chex.ArrayDevice + dvolumetric_strain: chex.ArrayDevice + volumetric_strain_centroid: chex.ArrayDevice + state_vars: dict + + +def init_particle_state( + loc: ArrayLike, + material: _Material, + element_ids: ArrayLike, +): + """Initialize a container of particles. + + Parameters + ---------- + loc: ArrayLike + Location of the particles. Expected shape (nparticles, 1, ndim) + material: diffmpm.materials._Material + Type of material for the set of particles. + element_ids: ArrayLike + The element ids that the particles belong to. This contains + information that will make sense only with the information of + the mesh that is being considered. + 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. + """ + loc = jnp.asarray(loc, dtype=jnp.float32) + if loc.ndim != 3: + raise ValueError( + f"`loc` should be of size (nparticles, 1, ndim); " f"found {loc.shape}" + ) + + mass = jnp.ones((loc.shape[0], 1, 1)) + density = jnp.ones_like(mass) * material.density + volume = jnp.ones_like(mass) + size = jnp.zeros_like(loc) + velocity = jnp.zeros_like(loc) + acceleration = jnp.zeros_like(loc) + momentum = jnp.zeros_like(loc) + strain = jnp.zeros((loc.shape[0], 6, 1)) + stress = jnp.zeros((loc.shape[0], 6, 1)) + strain_rate = jnp.zeros((loc.shape[0], 6, 1)) + dstrain = jnp.zeros((loc.shape[0], 6, 1)) + f_ext = jnp.zeros_like(loc) + traction = jnp.zeros_like(loc) + reference_loc = jnp.zeros_like(loc) + dvolumetric_strain = jnp.zeros((loc.shape[0], 1)) + volumetric_strain_centroid = jnp.zeros((loc.shape[0], 1)) + state_vars = {} + if material.state_vars: + state_vars = material.initialize_state_variables(loc.shape[0]) + return _ParticlesState( + nparticles=loc.shape[0], + loc=loc, + material=material, + element_ids=element_ids, + mass=mass, + density=density, + volume=volume, + size=size, + velocity=velocity, + acceleration=acceleration, + momentum=momentum, + strain=strain, + stress=stress, + strain_rate=strain_rate, + dstrain=dstrain, + f_ext=f_ext, + traction=traction, + reference_loc=reference_loc, + dvolumetric_strain=dvolumetric_strain, + volumetric_strain_centroid=volumetric_strain_centroid, + state_vars=state_vars, + ) + + +# TODO: Can these methods just return the updated arrays to +# a single function which then generates the new state? + + +def set_mass_volume(state, m: ArrayLike) -> _ParticlesState: + """Set particle mass. + + Parameters + ---------- + m: float, array_like + Mass to be set for particles. If scalar, mass for all + particles is set to this value. + """ + m = jnp.asarray(m) + if jnp.isscalar(m): + mass = jnp.ones_like(state.loc) * m + elif m.shape == state.mass.shape: + mass = m + else: + raise ValueError( + f"Incompatible shapes. Expected {state.mass.shape}, " f"found {m.shape}." + ) + volume = jnp.divide(mass, state.material.properties["density"]) + return state.replace(mass=mass, volume=volume) + + +def compute_volume(state, elements: _Element, total_elements: int): + """Compute volume of all particles. + + Parameters + ---------- + state: + Current state + elements: diffmpm._Element + Elements that the particles are present in, and are used to + compute the particles' volumes. + total_elements: int + Total elements present in `elements`. + """ + particles_per_element = jnp.bincount( + state.element_ids, length=elements.total_elements + ) + vol = ( + elements.volume.squeeze((1, 2))[state.element_ids] # type: ignore + / particles_per_element[state.element_ids] + ) + volume = state.volume.at[:, 0, 0].set(vol) + size = state.size.at[:].set(volume ** (1 / state.size.shape[-1])) + mass = state.mass.at[:, 0, 0].set(vol * state.density.squeeze()) + return state.replace(mass=mass, size=size, volume=volume) + + +def update_natural_coords(state, elements: _Element): + r"""Update natural coordinates for the particles. + + Whenever the particles' physical coordinates change, their + natural coordinates need to be updated. This function updates + the natural coordinates of the particles based on the element + a particle is a part of. The update formula is + + \[ + \xi = (2x - (x_1^e + x_2^e)) / (x_2^e - x_1^e) + \] + + where \(x_i^e\) are the nodal coordinates of the element the + particle is in. If a particle is not in any element + (element_id = -1), its natural coordinate is set to 0. + + Parameters + ---------- + elements: diffmpm.element._Element + Elements based on which to update the natural coordinates + of the particles. + """ + t = vmap(jit(elements.id_to_node_loc))(state.element_ids) + xi_coords = (state.loc - (t[:, 0, ...] + t[:, 2, ...]) / 2) * ( + 2 / (t[:, 2, ...] - t[:, 0, ...]) + ) + return state.replace(reference_loc=xi_coords) + + +def update_position_velocity( + state, elements: _Element, dt: float, velocity_update: bool +): + """Transfer nodal velocity to particles and update particle position. + + The velocity is calculated based on the total force at nodes. + + Parameters + ---------- + elements: diffmpm.element._Element + Elements whose nodes are used to transfer the velocity. + dt: float + Timestep. + velocity_update: bool + If True, velocity is directly used as nodal velocity, else + velocity is calculated is interpolated nodal acceleration + multiplied by dt. Default is False. + """ + mapped_positions = elements.shapefn(state.reference_loc) + mapped_ids = vmap(jit(elements.id_to_node_ids))(state.element_ids).squeeze(-1) + nodal_velocity = jnp.sum( + mapped_positions * elements.nodes.velocity[mapped_ids], axis=1 + ) + nodal_acceleration = jnp.sum( + mapped_positions * elements.nodes.acceleration[mapped_ids], + axis=1, + ) + velocity = state.velocity.at[:].set( + lax.cond( + velocity_update, + lambda sv, nv, na, t: nv, + lambda sv, nv, na, t: sv + na * t, + state.velocity, + nodal_velocity, + nodal_acceleration, + dt, + ) + ) + loc = state.loc.at[:].add(nodal_velocity * dt) + momentum = state.momentum.at[:].set(state.mass * state.velocity) + return state.replace(velocity=velocity, loc=loc, momentum=momentum) + + +def _compute_strain_rate(state, dn_dx: ArrayLike, elements: _Element): + """Compute the strain rate for particles. + + Parameters + ---------- + dn_dx: ArrayLike + The gradient of the shape function. Expected shape + `(nparticles, 1, ndim)` + elements: diffmpm.element._Element + Elements whose nodes are used to calculate the strain rate. + """ + dn_dx = jnp.asarray(dn_dx) + strain_rate = jnp.zeros((dn_dx.shape[0], 6, 1)) # (nparticles, 6, 1) + mapped_vel = vmap(jit(elements.id_to_node_vel))( + state.element_ids + ) # (nparticles, 2, 1) + + temp = mapped_vel.squeeze(2) + + @jit + def _step(pid, args): + dndx, nvel, strain_rate = args + matmul = dndx[pid].T @ nvel[pid] + strain_rate = strain_rate.at[pid, 0].add(matmul[0, 0]) + strain_rate = strain_rate.at[pid, 1].add(matmul[1, 1]) + strain_rate = strain_rate.at[pid, 3].add(matmul[0, 1] + matmul[1, 0]) + return dndx, nvel, strain_rate + + args = (dn_dx, temp, strain_rate) + _, _, strain_rate = lax.fori_loop(0, state.loc.shape[0], _step, args) + strain_rate = jnp.where( + jnp.abs(strain_rate) < 1e-12, jnp.zeros_like(strain_rate), strain_rate + ) + return strain_rate + + +def compute_strain(state, elements: _Element, dt: float): + """Compute the strain on all particles. + + This is done by first calculating the strain rate for the particles + and then calculating strain as `strain += strain rate * dt`. + + Parameters + ---------- + elements: diffmpm.element._Element + Elements whose nodes are used to calculate the strain. + dt : float + Timestep. + """ + # breakpoint() + mapped_coords = vmap(jit(elements.id_to_node_loc))(state.element_ids).squeeze(2) + dn_dx_ = vmap(jit(elements.shapefn_grad))( + state.reference_loc[:, jnp.newaxis, ...], mapped_coords + ) + strain_rate = _compute_strain_rate(state, dn_dx_, elements) + dstrain = state.dstrain.at[:].set(strain_rate * dt) + + strain = state.strain.at[:].add(dstrain) + centroids = jnp.zeros_like(state.loc) + dn_dx_centroid_ = vmap(jit(elements.shapefn_grad))( + centroids[:, jnp.newaxis, ...], mapped_coords + ) + strain_rate_centroid = _compute_strain_rate(state, dn_dx_centroid_, elements) + ndim = state.loc.shape[-1] + dvolumetric_strain = dt * strain_rate_centroid[:, :ndim].sum(axis=1) + volumetric_strain_centroid = state.volumetric_strain_centroid.at[:].add( + dvolumetric_strain + ) + return state.replace( + strain_rate=strain_rate, + dstrain=dstrain, + strain=strain, + dvolumetric_strain=dvolumetric_strain, + volumetric_strain_centroid=volumetric_strain_centroid, + ) + + +def compute_stress(state, *args): + """Compute the strain on all particles. + + This calculation is governed by the material of the + particles. The stress calculated by the material is then + added to the particles current stress values. + """ + stress = state.stress.at[:].add(state.material.compute_stress(state)) + return state.replace(stress=stress) + + +def update_volume(state, *args): + """Update volume based on central strain rate.""" + volume = state.volume.at[:, 0, :].multiply(1 + state.dvolumetric_strain) + density = state.density.at[:, 0, :].divide(1 + state.dvolumetric_strain) + return state.replace(volume=volume, density=density) + + +def assign_traction(state, pids: ArrayLike, dir: int, traction_: float): + """Assign traction to particles. + + Parameters + ---------- + pids: ArrayLike + IDs of the particles to which traction should be applied. + dir: int + The direction in which traction should be applied. + traction_: float + Traction value to be applied in the direction. + """ + traction = state.traction.at[pids, 0, dir].add( + traction_ * state.volume[pids, 0, 0] / state.size[pids, 0, dir] + ) + return traction + + +def zero_traction(state, *args): + """Set all traction values to 0.""" + traction = state.traction.at[:].set(0) + return state.replace(traction=traction) + @register_pytree_node_class class Particles(Sized): diff --git a/tests/test_newtonian.py b/tests/newtonian.py similarity index 100% rename from tests/test_newtonian.py rename to tests/newtonian.py diff --git a/tests/test_element.py b/tests/test_element.py index 11e9d6a..553693b 100644 --- a/tests/test_element.py +++ b/tests/test_element.py @@ -5,8 +5,9 @@ from diffmpm.element import Quadrilateral4Node from diffmpm.forces import NodalForce from diffmpm.functions import Unit -from diffmpm.materials import SimpleMaterial -from diffmpm.particle import Particles +from diffmpm.materials import init_simple +from diffmpm.particle import init_particle_state +import diffmpm.particle as dpar class TestLinear1D: @@ -21,8 +22,8 @@ def elements(self): @pytest.fixture def particles(self): loc = jnp.array([[0.5, 0.5], [0.5, 0.5]]).reshape(2, 1, 2) - material = SimpleMaterial({"E": 1, "density": 1}) - return Particles(loc, material, jnp.array([0, 0])) + material = init_simple({"id": 0, "E": 1, "density": 1}) + return init_particle_state(loc, material, jnp.array([0, 0])) @pytest.mark.parametrize( "particle_coords, expected", @@ -107,19 +108,20 @@ def test_element_node_vel(self, elements): assert jnp.all(node_vel == true_vel) def test_compute_nodal_mass(self, elements, particles): + # particles = particles.replace(mass=particles.mass + 1) particles.mass += 1 elements.compute_nodal_mass(particles) true_mass = jnp.ones((4, 1, 1)) assert jnp.all(elements.nodes.mass == true_mass) def test_compute_nodal_momentum(self, elements, particles): - particles.velocity += 1 + particles = particles.replace(velocity=particles.velocity + 1) elements.compute_nodal_momentum(particles) true_momentum = jnp.ones((4, 1, 1)) * 0.5 assert jnp.all(elements.nodes.momentum == true_momentum) def test_compute_external_force(self, elements, particles): - particles.f_ext += 1 + particles = particles.replace(f_ext=particles.f_ext + 1) elements.compute_external_force(particles) true_fext = jnp.ones((4, 1, 1)) * 0.5 assert jnp.all(elements.nodes.f_ext == true_fext) @@ -135,7 +137,7 @@ def test_compute_external_force(self, elements, particles): ], ) def test_compute_body_force(self, elements, particles, gravity, expected): - particles.mass += 1 + particles = particles.replace(mass=particles.mass + 1) elements.compute_body_force(particles, gravity) assert jnp.all(elements.nodes.f_ext == expected) @@ -194,12 +196,12 @@ def test_update_nodal_acceleration_velocity(self, elements, particles): ) def test_set_particle_element_ids(self, elements, particles): - particles.element_ids = jnp.array([-1, -1]) + particles = particles.replace(element_ids=jnp.array([-1, -1])) elements.set_particle_element_ids(particles) assert jnp.all(particles.element_ids == jnp.array([0, 0])) def test_compute_internal_force(self, elements, particles): - particles.compute_volume(elements, 1) + particles = dpar.compute_volume(particles, elements, 1) particles.stress += 1 elements.compute_internal_force(particles) assert jnp.allclose( @@ -212,7 +214,7 @@ def test_compute_volume(self, elements): assert jnp.allclose(elements.volume, jnp.array([1]).reshape(1, 1, 1)) def test_apply_particle_traction_forces(self, elements, particles): - particles.traction += jnp.array([1, 0]) + particles = particles.replace(traction=particles.traction + jnp.array([1, 0])) elements.apply_particle_traction_forces(particles) assert jnp.allclose( elements.nodes.f_ext, diff --git a/tests/test_material.py b/tests/test_material.py index f81cfb0..aef7ee4 100644 --- a/tests/test_material.py +++ b/tests/test_material.py @@ -1,32 +1,34 @@ import jax.numpy as jnp import pytest -from diffmpm.materials import LinearElastic, SimpleMaterial -from diffmpm.particle import Particles +from diffmpm.materials import init_linear_elastic, init_simple +from diffmpm.particle import init_particle_state particles_dstrain_stress_targets = [ ( - Particles( + init_particle_state( jnp.array([[0.5, 0.5]]).reshape(1, 1, 2), - SimpleMaterial({"E": 10, "density": 1}), + init_simple({"id": 0, "E": 10, "density": 1}), jnp.array([0]), ), jnp.ones((1, 6, 1)), jnp.ones((1, 6, 1)) * 10, ), ( - Particles( + init_particle_state( jnp.array([[0.5, 0.5]]).reshape(1, 1, 2), - LinearElastic({"density": 1, "youngs_modulus": 10, "poisson_ratio": 1}), + init_linear_elastic( + {"id": 0, "density": 1, "youngs_modulus": 10, "poisson_ratio": 1} + ), jnp.array([0]), ), jnp.ones((1, 6, 1)), jnp.array([-10, -10, -10, 2.5, 2.5, 2.5]).reshape(1, 6, 1), ), ( - Particles( + init_particle_state( jnp.array([[0.5, 0.5]]).reshape(1, 1, 2), - LinearElastic( - {"density": 1000, "youngs_modulus": 1e7, "poisson_ratio": 0.3} + init_linear_elastic( + {"id": 0, "density": 1000, "youngs_modulus": 1e7, "poisson_ratio": 0.3} ), jnp.array([0]), ), @@ -36,10 +38,10 @@ ), ), ( - Particles( + init_particle_state( jnp.array([[0.5, 0.5]]).reshape(1, 1, 2), - LinearElastic( - {"density": 1000, "youngs_modulus": 1e7, "poisson_ratio": 0.3} + init_linear_elastic( + {"id": 0, "density": 1000, "youngs_modulus": 1e7, "poisson_ratio": 0.3} ), jnp.array([0]), ), diff --git a/tests/test_particle.py b/tests/test_particle.py index 07b8c19..9c2acae 100644 --- a/tests/test_particle.py +++ b/tests/test_particle.py @@ -2,8 +2,9 @@ import pytest from diffmpm.element import Quadrilateral4Node -from diffmpm.materials import SimpleMaterial -from diffmpm.particle import Particles +from diffmpm.materials import init_simple +from diffmpm.particle import init_particle_state +import diffmpm.particle as dpar class TestParticles: @@ -14,8 +15,8 @@ def elements(self): @pytest.fixture def particles(self): loc = jnp.array([[0.5, 0.5], [0.5, 0.5]]).reshape(2, 1, 2) - material = SimpleMaterial({"E": 1, "density": 1}) - return Particles(loc, material, jnp.array([0, 0])) + material = init_simple({"id": 0, "E": 1, "density": 1}) + return init_particle_state(loc, material, jnp.array([0, 0])) @pytest.mark.parametrize( "velocity_update, expected", @@ -25,20 +26,22 @@ def particles(self): ], ) def test_update_velocity(self, elements, particles, velocity_update, expected): - particles.update_natural_coords(elements) + dpar.update_natural_coords(particles, elements) elements.nodes = elements.nodes.replace( acceleration=elements.nodes.acceleration + 1 ) elements.nodes = elements.nodes.replace(velocity=elements.nodes.velocity + 1) - particles.update_position_velocity(elements, 0.1, velocity_update) + particles = dpar.update_position_velocity( + particles, elements, 0.1, velocity_update + ) assert jnp.allclose(particles.velocity, expected) def test_compute_strain(self, elements, particles): elements.nodes = elements.nodes.replace( velocity=jnp.array([[0, 1], [0, 2], [0, 3], [0, 4]]).reshape(4, 1, 2) ) - particles.update_natural_coords(elements) - particles.compute_strain(elements, 0.1) + particles = dpar.update_natural_coords(particles, elements) + particles = dpar.compute_strain(particles, elements, 0.1) assert jnp.allclose( particles.strain, jnp.array([[0, 0.2, 0, 0.1, 0, 0], [0, 0.2, 0, 0.1, 0, 0]]).reshape( @@ -48,17 +51,17 @@ def test_compute_strain(self, elements, particles): assert jnp.allclose(particles.volumetric_strain_centroid, jnp.array([0.2])) def test_compute_volume(self, elements, particles): - particles.compute_volume(elements, elements.total_elements) + particles = dpar.compute_volume(particles, elements, elements.total_elements) assert jnp.allclose(particles.volume, jnp.array([0.5, 0.5]).reshape(2, 1, 1)) def test_assign_traction(self, elements, particles): - particles.compute_volume(elements, elements.total_elements) - particles.assign_traction(jnp.array([0]), 1, 10) + particles = dpar.compute_volume(particles, elements, elements.total_elements) + traction = dpar.assign_traction(particles, jnp.array([0]), 1, 10) assert jnp.allclose( - particles.traction, jnp.array([[0, 7.071068], [0, 0]]).reshape(2, 1, 2) + traction, jnp.array([[0, 7.071068], [0, 0]]).reshape(2, 1, 2) ) def test_zero_traction(self, particles): - particles.traction += 1 - particles.zero_traction() + particles = particles.replace(traction=particles.traction + 1) + particles = dpar.zero_traction(particles) assert jnp.all(particles.traction == 0) From 4158e11483da37c1c4af05355b343f0c969a2b5d Mon Sep 17 00:00:00 2001 From: Chahak Mehta Date: Sat, 22 Jul 2023 21:31:02 -0700 Subject: [PATCH 12/21] Working frozen particle state --- diffmpm/constraint.py | 2 +- diffmpm/element.py | 36 +++-- diffmpm/io.py | 7 +- diffmpm/mesh.py | 8 +- diffmpm/particle.py | 355 +---------------------------------------- tests/test_element.py | 8 +- tests/test_material.py | 2 +- 7 files changed, 39 insertions(+), 379 deletions(-) diff --git a/diffmpm/constraint.py b/diffmpm/constraint.py index 2f8d7b5..dfe5bb7 100644 --- a/diffmpm/constraint.py +++ b/diffmpm/constraint.py @@ -31,7 +31,7 @@ def apply(self, obj, ids): Parameters ---------- - obj : diffmpm.node.Nodes, diffmpm.particle.Particles + obj : diffmpm.node.Nodes, diffmpm.particle._ParticlesState Object on which the constraint is applied ids : array_like The indices of the container `obj` on which the constraint diff --git a/diffmpm/element.py b/diffmpm/element.py index bf3d0d6..b082abb 100644 --- a/diffmpm/element.py +++ b/diffmpm/element.py @@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Optional, Sequence, Tuple if TYPE_CHECKING: - from diffmpm.particle import Particles + from diffmpm.particle import _ParticlesState import jax.numpy as jnp from jax import Array, jacobian, jit, lax, vmap @@ -117,12 +117,12 @@ def shapefn_grad(self, xi: ArrayLike, coords: ArrayLike): ... @abc.abstractmethod - def set_particle_element_ids(self, particles: Particles): + def set_particle_element_ids(self, particles: _ParticlesState): """Set the element IDs that particles are present in.""" ... # Mapping from particles to nodes (P2G) - def compute_nodal_mass(self, particles: Particles): + def compute_nodal_mass(self, particles: _ParticlesState): r"""Compute the nodal mass based on particle mass. The nodal mass is updated as a sum of particle mass for @@ -157,7 +157,7 @@ def _step(pid, args): # TODO: Return state instead of setting self.nodes = self.nodes.replace(mass=mass) - def compute_nodal_momentum(self, particles: Particles): + def compute_nodal_momentum(self, particles: _ParticlesState): r"""Compute the nodal mass based on particle mass. The nodal mass is updated as a sum of particle mass for @@ -193,7 +193,7 @@ def _step(pid, args): # TODO: Return state instead of setting self.nodes = self.nodes.replace(momentum=new_momentum) - def compute_velocity(self, particles: Particles): + def compute_velocity(self, particles: _ParticlesState): """Compute velocity using momentum.""" velocity = jnp.where( self.nodes.mass == 0, @@ -208,7 +208,7 @@ def compute_velocity(self, particles: Particles): # TODO: Return state instead of setting self.nodes = self.nodes.replace(velocity=velocity) - def compute_external_force(self, particles: Particles): + def compute_external_force(self, particles: _ParticlesState): r"""Update the nodal external force based on particle f_ext. The nodal force is updated as a sum of particle external @@ -243,7 +243,7 @@ def _step(pid, args): # TODO: Return state instead of setting self.nodes = self.nodes.replace(f_ext=f_ext) - def compute_body_force(self, particles: Particles, gravity: ArrayLike): + def compute_body_force(self, particles: _ParticlesState, gravity: ArrayLike): r"""Update the nodal external force based on particle mass. The nodal force is updated as a sum of particle body @@ -255,7 +255,7 @@ def compute_body_force(self, particles: Particles, gravity: ArrayLike): Parameters ---------- - particles: diffmpm.particle.Particles + particles: diffmpm.particle._ParticlesState Particles to map to the nodal values. """ @@ -280,12 +280,14 @@ def _step(pid, args): # TODO: Return state instead of setting self.nodes = self.nodes.replace(f_ext=f_ext) - def apply_concentrated_nodal_forces(self, particles: Particles, curr_time: float): + def apply_concentrated_nodal_forces( + self, particles: _ParticlesState, curr_time: float + ): """Apply concentrated nodal forces. Parameters ---------- - particles: Particles + particles: _ParticlesState Particles in the simulation. curr_time: float Current time in the simulation. @@ -315,7 +317,7 @@ def _f(x, *, orig): # TODO: Return state instead of setting self.nodes = self.nodes.replace(f_ext=f_ext) - def apply_particle_traction_forces(self, particles: Particles): + def apply_particle_traction_forces(self, particles: _ParticlesState): """Apply concentrated nodal forces. Parameters @@ -338,7 +340,7 @@ def _step(pid, args): self.nodes = self.nodes.replace(f_ext=f_ext) def update_nodal_acceleration_velocity( - self, particles: Particles, dt: float, *args + self, particles: _ParticlesState, dt: float, *args ): """Update the nodal momentum based on total force on nodes.""" total_force = self.nodes.f_int + self.nodes.f_ext + self.nodes.f_damp @@ -631,7 +633,7 @@ def compute_internal_force(self, particles): Parameters ---------- - particles: diffmpm.particle.Particles + particles: diffmpm.particle._ParticlesState Particles to map to the nodal values. """ @@ -891,7 +893,7 @@ def shapefn_grad(self, xi: ArrayLike, coords: ArrayLike): result = grad_sf @ jnp.linalg.inv(_jacobian).T return result - def set_particle_element_ids(self, particles: Particles): + def set_particle_element_ids(self, particles: _ParticlesState): """Set the element IDs for the particles. If the particle doesn't lie between the boundaries of any @@ -911,9 +913,9 @@ def f(x): return element_id ids = vmap(f)(particles.loc) - particles.element_ids = ids + return particles.replace(element_ids=ids) - def compute_internal_force(self, particles: Particles): + def compute_internal_force(self, particles: _ParticlesState): r"""Update the nodal internal force based on particle mass. The nodal force is updated as a sum of internal forces for @@ -927,7 +929,7 @@ def compute_internal_force(self, particles: Particles): Parameters ---------- - particles: diffmpm.particle.Particles + particles: diffmpm.particle._ParticlesState Particles to map to the nodal values. """ diff --git a/diffmpm/io.py b/diffmpm/io.py index 6a4f46e..822e50c 100644 --- a/diffmpm/io.py +++ b/diffmpm/io.py @@ -9,7 +9,7 @@ from diffmpm.constraint import Constraint from diffmpm.forces import NodalForce, ParticleTraction from diffmpm.functions import Linear, Unit -from diffmpm.particle import Particles, _ParticlesState, init_particle_state +from diffmpm.particle import _ParticlesState, init_particle_state class Config: @@ -71,8 +71,9 @@ def _parse_particles(self, config): with open(pset_config["file"], "r") as f: ploc = jnp.asarray(json.load(f)) peids = jnp.zeros(ploc.shape[0], dtype=jnp.int32) - pset = init_particle_state(ploc, pmat, peids) - pset.velocity = pset.velocity.at[:].set(pset_config["init_velocity"]) + pset = init_particle_state( + ploc, pmat, peids, init_vel=jnp.asarray(pset_config["init_velocity"]) + ) particle_sets.append(pset) self.parsed_config["particles"] = particle_sets diff --git a/diffmpm/mesh.py b/diffmpm/mesh.py index ce68d6e..15e4fe2 100644 --- a/diffmpm/mesh.py +++ b/diffmpm/mesh.py @@ -7,7 +7,7 @@ from jax.tree_util import register_pytree_node_class, tree_map from diffmpm.element import _Element -from diffmpm.particle import Particles, _ParticlesState +from diffmpm.particle import _ParticlesState import diffmpm.particle as dpart from diffmpm.forces import ParticleTraction @@ -45,14 +45,16 @@ def apply_on_elements(self, function: str, args: Tuple = ()): f = getattr(self.elements, function) def _func(particles, *, func, fargs): - func(particles, *fargs) + return func(particles, *fargs) partial_func = partial(_func, func=f, fargs=args) - tree_map( + _out = tree_map( partial_func, self.particles, is_leaf=lambda x: isinstance(x, _ParticlesState), ) + if function == "set_particle_element_ids": + self.particles = _out # TODO: Convert to using jax directives for loop def apply_on_particles(self, function: str, args: Tuple = ()): diff --git a/diffmpm/particle.py b/diffmpm/particle.py index bbe17a2..97921e0 100644 --- a/diffmpm/particle.py +++ b/diffmpm/particle.py @@ -11,7 +11,7 @@ import chex -@chex.dataclass() +@chex.dataclass(frozen=True) class _ParticlesState: nparticles: int loc: chex.ArrayDevice @@ -37,9 +37,10 @@ class _ParticlesState: def init_particle_state( - loc: ArrayLike, + loc: chex.ArrayDevice, material: _Material, - element_ids: ArrayLike, + element_ids: chex.ArrayDevice, + init_vel: chex.ArrayDevice = 0, ): """Initialize a container of particles. @@ -71,7 +72,7 @@ def init_particle_state( density = jnp.ones_like(mass) * material.density volume = jnp.ones_like(mass) size = jnp.zeros_like(loc) - velocity = jnp.zeros_like(loc) + velocity = jnp.ones_like(loc) * init_vel acceleration = jnp.zeros_like(loc) momentum = jnp.zeros_like(loc) strain = jnp.zeros((loc.shape[0], 6, 1)) @@ -352,349 +353,3 @@ def zero_traction(state, *args): """Set all traction values to 0.""" traction = state.traction.at[:].set(0) return state.replace(traction=traction) - - -@register_pytree_node_class -class Particles(Sized): - """Container class for a set of particles.""" - - def __init__( - self, - loc: ArrayLike, - material: _Material, - element_ids: ArrayLike, - initialized: Optional[bool] = None, - data: Optional[Tuple[ArrayLike, ...]] = None, - ): - """Initialize a container of particles. - - Parameters - ---------- - loc: ArrayLike - Location of the particles. Expected shape (nparticles, 1, ndim) - material: diffmpm.materials._Material - Type of material for the set of particles. - element_ids: ArrayLike - The element ids that the particles belong to. This contains - information that will make sense only with the information of - the mesh that is being considered. - 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. - """ - self.material = material - self.element_ids = element_ids - loc = jnp.asarray(loc, dtype=jnp.float32) - if loc.ndim != 3: - raise ValueError( - f"`loc` should be of size (nparticles, 1, ndim); " f"found {loc.shape}" - ) - self.loc = loc - - if initialized is None: - self.mass = jnp.ones((self.loc.shape[0], 1, 1)) - self.density = ( - jnp.ones_like(self.mass) * self.material.properties["density"] - ) - self.volume = jnp.ones_like(self.mass) - self.size = jnp.zeros_like(self.loc) - self.velocity = jnp.zeros_like(self.loc) - self.acceleration = jnp.zeros_like(self.loc) - self.momentum = jnp.zeros_like(self.loc) - self.strain = jnp.zeros((self.loc.shape[0], 6, 1)) - self.stress = jnp.zeros((self.loc.shape[0], 6, 1)) - self.strain_rate = jnp.zeros((self.loc.shape[0], 6, 1)) - self.dstrain = jnp.zeros((self.loc.shape[0], 6, 1)) - self.f_ext = jnp.zeros_like(self.loc) - self.traction = jnp.zeros_like(self.loc) - self.reference_loc = jnp.zeros_like(self.loc) - self.dvolumetric_strain = jnp.zeros((self.loc.shape[0], 1)) - self.volumetric_strain_centroid = jnp.zeros((self.loc.shape[0], 1)) - self.state_vars = {} - if self.material.state_vars: - self.state_vars = self.material.initialize_state_variables( - self.loc.shape[0] - ) - else: - ( - self.mass, - self.density, - self.volume, - self.size, - self.velocity, - self.acceleration, - self.momentum, - self.strain, - self.stress, - self.strain_rate, - self.dstrain, - self.f_ext, - self.traction, - self.reference_loc, - self.dvolumetric_strain, - self.volumetric_strain_centroid, - self.state_vars, - ) = data # type: ignore - self.initialized = True - - def tree_flatten(self): - """Flatten class as Pytree type.""" - children = ( - self.loc, - self.element_ids, - self.initialized, - self.mass, - self.density, - self.volume, - self.size, - self.velocity, - self.acceleration, - self.momentum, - self.strain, - self.stress, - self.strain_rate, - self.dstrain, - self.f_ext, - self.traction, - self.reference_loc, - self.dvolumetric_strain, - self.volumetric_strain_centroid, - self.state_vars, - ) - aux_data = (self.material,) - return (children, aux_data) - - @classmethod - def tree_unflatten(cls, aux_data, children): - """Unflatten class from Pytree type.""" - return cls( - children[0], - aux_data[0], - children[1], - initialized=children[2], - data=children[3:], - ) - - def __len__(self) -> int: - """Set length of the class as number of particles.""" - return self.loc.shape[0] - - def __repr__(self) -> str: - """Informative repr showing number of particles.""" - return f"Particles(nparticles={len(self)})" - - def set_mass_volume(self, m: ArrayLike): - """Set particle mass. - - Parameters - ---------- - m: float, array_like - Mass to be set for particles. If scalar, mass for all - particles is set to this value. - """ - m = jnp.asarray(m) - if jnp.isscalar(m): - self.mass = jnp.ones_like(self.loc) * m - elif m.shape == self.mass.shape: - self.mass = m - else: - raise ValueError( - f"Incompatible shapes. Expected {self.mass.shape}, " f"found {m.shape}." - ) - self.volume = jnp.divide(self.mass, self.material.properties["density"]) - - def compute_volume(self, elements: _Element, total_elements: int): - """Compute volume of all particles. - - Parameters - ---------- - elements: diffmpm._Element - Elements that the particles are present in, and are used to - compute the particles' volumes. - total_elements: int - Total elements present in `elements`. - """ - particles_per_element = jnp.bincount( - self.element_ids, length=elements.total_elements - ) - vol = ( - elements.volume.squeeze((1, 2))[self.element_ids] # type: ignore - / particles_per_element[self.element_ids] - ) - self.volume = self.volume.at[:, 0, 0].set(vol) - self.size = self.size.at[:].set(self.volume ** (1 / self.size.shape[-1])) - self.mass = self.mass.at[:, 0, 0].set(vol * self.density.squeeze()) - - def update_natural_coords(self, elements: _Element): - r"""Update natural coordinates for the particles. - - Whenever the particles' physical coordinates change, their - natural coordinates need to be updated. This function updates - the natural coordinates of the particles based on the element - a particle is a part of. The update formula is - - \[ - \xi = (2x - (x_1^e + x_2^e)) / (x_2^e - x_1^e) - \] - - where \(x_i^e\) are the nodal coordinates of the element the - particle is in. If a particle is not in any element - (element_id = -1), its natural coordinate is set to 0. - - Parameters - ---------- - elements: diffmpm.element._Element - Elements based on which to update the natural coordinates - of the particles. - """ - t = vmap(jit(elements.id_to_node_loc))(self.element_ids) - xi_coords = (self.loc - (t[:, 0, ...] + t[:, 2, ...]) / 2) * ( - 2 / (t[:, 2, ...] - t[:, 0, ...]) - ) - self.reference_loc = xi_coords - - def update_position_velocity( - self, elements: _Element, dt: float, velocity_update: bool - ): - """Transfer nodal velocity to particles and update particle position. - - The velocity is calculated based on the total force at nodes. - - Parameters - ---------- - elements: diffmpm.element._Element - Elements whose nodes are used to transfer the velocity. - dt: float - Timestep. - velocity_update: bool - If True, velocity is directly used as nodal velocity, else - velocity is calculated is interpolated nodal acceleration - multiplied by dt. Default is False. - """ - mapped_positions = elements.shapefn(self.reference_loc) - mapped_ids = vmap(jit(elements.id_to_node_ids))(self.element_ids).squeeze(-1) - nodal_velocity = jnp.sum( - mapped_positions * elements.nodes.velocity[mapped_ids], axis=1 - ) - nodal_acceleration = jnp.sum( - mapped_positions * elements.nodes.acceleration[mapped_ids], - axis=1, - ) - self.velocity = self.velocity.at[:].set( - lax.cond( - velocity_update, - lambda sv, nv, na, t: nv, - lambda sv, nv, na, t: sv + na * t, - self.velocity, - nodal_velocity, - nodal_acceleration, - dt, - ) - ) - self.loc = self.loc.at[:].add(nodal_velocity * dt) - self.momentum = self.momentum.at[:].set(self.mass * self.velocity) - - def compute_strain(self, elements: _Element, dt: float): - """Compute the strain on all particles. - - This is done by first calculating the strain rate for the particles - and then calculating strain as `strain += strain rate * dt`. - - Parameters - ---------- - elements: diffmpm.element._Element - Elements whose nodes are used to calculate the strain. - dt : float - Timestep. - """ - mapped_coords = vmap(jit(elements.id_to_node_loc))(self.element_ids).squeeze(2) - dn_dx_ = vmap(jit(elements.shapefn_grad))( - self.reference_loc[:, jnp.newaxis, ...], mapped_coords - ) - self.strain_rate = self._compute_strain_rate(dn_dx_, elements) - self.dstrain = self.dstrain.at[:].set(self.strain_rate * dt) - - self.strain = self.strain.at[:].add(self.dstrain) - centroids = jnp.zeros_like(self.loc) - dn_dx_centroid_ = vmap(jit(elements.shapefn_grad))( - centroids[:, jnp.newaxis, ...], mapped_coords - ) - strain_rate_centroid = self._compute_strain_rate(dn_dx_centroid_, elements) - ndim = self.loc.shape[-1] - self.dvolumetric_strain = dt * strain_rate_centroid[:, :ndim].sum(axis=1) - self.volumetric_strain_centroid = self.volumetric_strain_centroid.at[:].add( - self.dvolumetric_strain - ) - - def _compute_strain_rate(self, dn_dx: ArrayLike, elements: _Element): - """Compute the strain rate for particles. - - Parameters - ---------- - dn_dx: ArrayLike - The gradient of the shape function. Expected shape - `(nparticles, 1, ndim)` - elements: diffmpm.element._Element - Elements whose nodes are used to calculate the strain rate. - """ - dn_dx = jnp.asarray(dn_dx) - strain_rate = jnp.zeros((dn_dx.shape[0], 6, 1)) # (nparticles, 6, 1) - mapped_vel = vmap(jit(elements.id_to_node_vel))( - self.element_ids - ) # (nparticles, 2, 1) - - temp = mapped_vel.squeeze(2) - - @jit - def _step(pid, args): - dndx, nvel, strain_rate = args - matmul = dndx[pid].T @ nvel[pid] - strain_rate = strain_rate.at[pid, 0].add(matmul[0, 0]) - strain_rate = strain_rate.at[pid, 1].add(matmul[1, 1]) - strain_rate = strain_rate.at[pid, 3].add(matmul[0, 1] + matmul[1, 0]) - return dndx, nvel, strain_rate - - args = (dn_dx, temp, strain_rate) - _, _, strain_rate = lax.fori_loop(0, self.loc.shape[0], _step, args) - strain_rate = jnp.where( - jnp.abs(strain_rate) < 1e-12, jnp.zeros_like(strain_rate), strain_rate - ) - return strain_rate - - def compute_stress(self, *args): - """Compute the strain on all particles. - - This calculation is governed by the material of the - particles. The stress calculated by the material is then - added to the particles current stress values. - """ - self.stress = self.stress.at[:].add(self.material.compute_stress(self)) - - def update_volume(self, *args): - """Update volume based on central strain rate.""" - self.volume = self.volume.at[:, 0, :].multiply(1 + self.dvolumetric_strain) - self.density = self.density.at[:, 0, :].divide(1 + self.dvolumetric_strain) - - def assign_traction(self, pids: ArrayLike, dir: int, traction_: float): - """Assign traction to particles. - - Parameters - ---------- - pids: ArrayLike - IDs of the particles to which traction should be applied. - dir: int - The direction in which traction should be applied. - traction_: float - Traction value to be applied in the direction. - """ - self.traction = self.traction.at[pids, 0, dir].add( - traction_ * self.volume[pids, 0, 0] / self.size[pids, 0, dir] - ) - - def zero_traction(self, *args): - """Set all traction values to 0.""" - self.traction = self.traction.at[:].set(0) diff --git a/tests/test_element.py b/tests/test_element.py index 553693b..5f78de5 100644 --- a/tests/test_element.py +++ b/tests/test_element.py @@ -108,8 +108,8 @@ def test_element_node_vel(self, elements): assert jnp.all(node_vel == true_vel) def test_compute_nodal_mass(self, elements, particles): - # particles = particles.replace(mass=particles.mass + 1) - particles.mass += 1 + particles = particles.replace(mass=particles.mass + 1) + # particles.mass += 1 elements.compute_nodal_mass(particles) true_mass = jnp.ones((4, 1, 1)) assert jnp.all(elements.nodes.mass == true_mass) @@ -197,12 +197,12 @@ def test_update_nodal_acceleration_velocity(self, elements, particles): def test_set_particle_element_ids(self, elements, particles): particles = particles.replace(element_ids=jnp.array([-1, -1])) - elements.set_particle_element_ids(particles) + particles = elements.set_particle_element_ids(particles) assert jnp.all(particles.element_ids == jnp.array([0, 0])) def test_compute_internal_force(self, elements, particles): particles = dpar.compute_volume(particles, elements, 1) - particles.stress += 1 + particles = particles.replace(stress=particles.stress + 1) elements.compute_internal_force(particles) assert jnp.allclose( elements.nodes.f_int, diff --git a/tests/test_material.py b/tests/test_material.py index aef7ee4..075b9f7 100644 --- a/tests/test_material.py +++ b/tests/test_material.py @@ -55,6 +55,6 @@ @pytest.mark.parametrize("particles, dstrain, target", particles_dstrain_stress_targets) def test_compute_stress(particles, dstrain, target): - particles.dstrain = dstrain + particles = particles.replace(dstrain=dstrain) stress = particles.material.compute_stress(particles) assert jnp.allclose(stress, target) From 3d0e43e0016e1a626c108ef60ec922653bff5293 Mon Sep 17 00:00:00 2001 From: Chahak Mehta Date: Sun, 23 Jul 2023 11:56:26 -0700 Subject: [PATCH 13/21] Reset node state in scheme and separate bc apply --- .../test_benchmark.py | 3 +- diffmpm/constraint.py | 34 +++- diffmpm/element.py | 160 +++++++++++++----- diffmpm/mesh.py | 2 +- diffmpm/node.py | 24 ++- diffmpm/scheme.py | 3 + tests/test_element.py | 1 - 7 files changed, 166 insertions(+), 61 deletions(-) diff --git a/benchmarks/2d/uniaxial_particle_traction/test_benchmark.py b/benchmarks/2d/uniaxial_particle_traction/test_benchmark.py index 1a53171..cec2a34 100644 --- a/benchmarks/2d/uniaxial_particle_traction/test_benchmark.py +++ b/benchmarks/2d/uniaxial_particle_traction/test_benchmark.py @@ -36,4 +36,5 @@ def test_benchmarks(): assert jnp.round(result["stress"][0, :, 0].max() - 0.9999997782938734, 5) == 0.0 -test_benchmarks() +if __name__ == "__main__": + test_benchmarks() diff --git a/diffmpm/constraint.py b/diffmpm/constraint.py index dfe5bb7..cb40a3e 100644 --- a/diffmpm/constraint.py +++ b/diffmpm/constraint.py @@ -26,7 +26,7 @@ def tree_unflatten(cls, aux_data, children): del children return cls(*aux_data) - def apply(self, obj, ids): + def apply_vel(self, obj, ids): """Apply constraint values to the passed object. Parameters @@ -38,12 +38,34 @@ def apply(self, obj, ids): will be applied. """ velocity = obj.velocity.at[ids, :, self.dir].set(self.velocity) + return velocity + + def apply_mom(self, obj, ids): + """Apply constraint values to the passed object. + + Parameters + ---------- + obj : diffmpm.node.Nodes, diffmpm.particle._ParticlesState + Object on which the constraint is applied + ids : array_like + The indices of the container `obj` on which the constraint + will be applied. + """ momentum = obj.momentum.at[ids, :, self.dir].set( obj.mass[ids, :, 0] * self.velocity ) - acceleration = obj.acceleration.at[ids, :, self.dir].set(0) - # return obj.replace( - # velocity=velocity, momentum=momentum, acceleration=acceleration - # ) + return momentum - return velocity, momentum, acceleration + def apply_acc(self, obj, ids): + """Apply constraint values to the passed object. + + Parameters + ---------- + obj : diffmpm.node.Nodes, diffmpm.particle._ParticlesState + Object on which the constraint is applied + ids : array_like + The indices of the container `obj` on which the constraint + will be applied. + """ + acceleration = obj.acceleration.at[ids, :, self.dir].set(0) + return acceleration diff --git a/diffmpm/element.py b/diffmpm/element.py index b082abb..3c15ec6 100644 --- a/diffmpm/element.py +++ b/diffmpm/element.py @@ -9,18 +9,26 @@ from diffmpm.particle import _ParticlesState import jax.numpy as jnp -from jax import Array, jacobian, jit, lax, vmap +from jax import Array, jacobian, jit, lax, tree_util, vmap from jax.tree_util import register_pytree_node_class, tree_map, tree_reduce -from jax import tree_util from jax.typing import ArrayLike from diffmpm.constraint import Constraint from diffmpm.forces import NodalForce -from diffmpm.node import _NodesState, init_state +from diffmpm.node import _NodesState, init_node_state +import chex __all__ = ["_Element", "Linear1D", "Quadrilateral4Node"] +@chex.dataclass() +class _ElementState: + nodes: _NodesState + total_elements: int + concentrated_nodal_forces: Sequence + volume: chex.ArrayDevice + + class _Element(abc.ABC): """Base element class that is inherited by all types of Elements.""" @@ -144,7 +152,8 @@ def _step(pid, args): mass = mass.at[el_nodes[pid]].add(pmass[pid] * mapped_pos[pid]) return pmass, mass, mapped_pos, el_nodes - mass = self.nodes.mass.at[:].set(0) + # mass = self.nodes.mass.at[:].set(0) + mass = self.nodes.mass mapped_positions = self.shapefn(particles.reference_loc) mapped_nodes = vmap(jit(self.id_to_node_ids))(particles.element_ids).squeeze(-1) args = ( @@ -179,7 +188,8 @@ def _step(pid, args): new_mom = mom.at[el_nodes[pid]].add(mapped_pos[pid] @ pmom[pid]) return pmom, new_mom, mapped_pos, el_nodes - curr_mom = self.nodes.momentum.at[:].set(0) + # curr_mom = self.nodes.momentum.at[:].set(0) + curr_mom = self.nodes.momentum mapped_positions = self.shapefn(particles.reference_loc) mapped_nodes = vmap(jit(self.id_to_node_ids))(particles.element_ids).squeeze(-1) args = ( @@ -230,7 +240,8 @@ def _step(pid, args): f_ext = f_ext.at[el_nodes[pid]].add(mapped_pos[pid] @ pf_ext[pid]) return f_ext, pf_ext, mapped_pos, el_nodes - f_ext = self.nodes.f_ext.at[:].set(0) + # f_ext = self.nodes.f_ext.at[:].set(0) + f_ext = self.nodes.f_ext mapped_positions = self.shapefn(particles.reference_loc) mapped_nodes = vmap(jit(self.id_to_node_ids))(particles.element_ids).squeeze(-1) args = ( @@ -366,52 +377,110 @@ def update_nodal_acceleration_velocity( velocity=velocity, acceleration=acceleration, momentum=momentum ) - def apply_boundary_constraints(self, *args): + def _apply_boundary_constraints_vel(self, *args): """Apply boundary conditions for nodal velocity.""" - if self.constraints: + # This assumes that the constraints don't have overlapping + # conditions. In case it does, only the first constraint will + # be applied. + def _func2(constraint, *, nodes): + return constraint[1].apply_vel(nodes, constraint[0]) - def _func2(constraint, *, nodes): - return constraint[1].apply(nodes, constraint[0]) + partial_func = partial(_func2, nodes=self.nodes) + _out = tree_map( + partial_func, self.constraints, is_leaf=lambda x: isinstance(x, tuple) + ) - partial_func = partial(_func2, nodes=self.nodes) - _out = tree_map( - partial_func, self.constraints, is_leaf=lambda x: isinstance(x, tuple) - ) + # _temp = tree_util.tree_transpose( + # tree_util.tree_structure([0 for e in _out]), + # tree_util.tree_structure(_out[0]), + # _out, + # ) - _temp = tree_util.tree_transpose( - tree_util.tree_structure([0 for e in _out]), - tree_util.tree_structure(_out[0]), - _out, - ) + def _f(x, *, orig): + return jnp.where(x == orig, jnp.nan, x) - def _f(x, *, orig): - return jnp.where(x == orig, jnp.nan, x) + _pf = partial(_f, orig=self.nodes.velocity) + _step_1 = tree_map(_pf, _out) + vel = tree_reduce( + lambda x, y: jnp.where(jnp.isnan(y), x, y), + [self.nodes.velocity, _step_1], + ) - _pf = partial(_f, orig=self.nodes.velocity) - _step_1 = tree_map(_pf, _temp[0]) - vel = tree_reduce( - lambda x, y: jnp.where(jnp.isnan(y), x, y), - [self.nodes.velocity, _step_1], - ) + # TODO: Return state instead of setting + self.nodes = self.nodes.replace(velocity=vel) - _pf = partial(_f, orig=self.nodes.momentum) - _step_1 = tree_map(_pf, _temp[1]) - mom = tree_reduce( - lambda x, y: jnp.where(jnp.isnan(y), x, y), - [self.nodes.momentum, _step_1], - ) + def _apply_boundary_constraints_mom(self, *args): + """Apply boundary conditions for nodal momentum.""" - _pf = partial(_f, orig=self.nodes.acceleration) - _step_1 = tree_map(_pf, _temp[2]) - acc = tree_reduce( - lambda x, y: jnp.where(jnp.isnan(y), x, y), - [self.nodes.acceleration, _step_1], - ) - # TODO: Return state instead of setting - self.nodes = self.nodes.replace( - velocity=vel, momentum=mom, acceleration=acc - ) + # This assumes that the constraints don't have overlapping + # conditions. In case it does, only the first constraint will + # be applied. + def _func2(constraint, *, nodes): + return constraint[1].apply_mom(nodes, constraint[0]) + + partial_func = partial(_func2, nodes=self.nodes) + _out = tree_map( + partial_func, self.constraints, is_leaf=lambda x: isinstance(x, tuple) + ) + + # _temp = tree_util.tree_transpose( + # tree_util.tree_structure([0 for e in _out]), + # tree_util.tree_structure(_out[0]), + # _out, + # ) + + def _f(x, *, orig): + return jnp.where(x == orig, jnp.nan, x) + + _pf = partial(_f, orig=self.nodes.momentum) + _step_1 = tree_map(_pf, _out) + mom = tree_reduce( + lambda x, y: jnp.where(jnp.isnan(y), x, y), + [self.nodes.momentum, _step_1], + ) + + # TODO: Return state instead of setting + self.nodes = self.nodes.replace(momentum=mom) + + def _apply_boundary_constraints_acc(self, *args): + """Apply boundary conditions for nodal acceleration.""" + + # This assumes that the constraints don't have overlapping + # conditions. In case it does, only the first constraint will + # be applied. + def _func2(constraint, *, nodes): + return constraint[1].apply_acc(nodes, constraint[0]) + + partial_func = partial(_func2, nodes=self.nodes) + _out = tree_map( + partial_func, self.constraints, is_leaf=lambda x: isinstance(x, tuple) + ) + + # _temp = tree_util.tree_transpose( + # tree_util.tree_structure([0 for e in _out]), + # tree_util.tree_structure(_out[0]), + # _out, + # ) + + def _f(x, *, orig): + return jnp.where(x == orig, jnp.nan, x) + + _pf = partial(_f, orig=self.nodes.acceleration) + _step_1 = tree_map(_pf, _out) + acc = tree_reduce( + lambda x, y: jnp.where(jnp.isnan(y), x, y), + [self.nodes.acceleration, _step_1], + ) + + # TODO: Return state instead of setting + self.nodes = self.nodes.replace(acceleration=acc) + + def apply_boundary_constraints(self, *args): + if self.constraints: + self._apply_boundary_constraints_vel(*args) + self._apply_boundary_constraints_mom(*args) + self._apply_boundary_constraints_acc(*args) @register_pytree_node_class @@ -753,7 +822,7 @@ def __init__( node_locations = ( jnp.asarray([coords[:, 1], coords[:, 0]]).T * self.el_len ).reshape(-1, 1, 2) - self.nodes = init_state(int(total_nodes), node_locations) + self.nodes = init_node_state(int(total_nodes), node_locations) else: self.nodes = nodes @@ -961,7 +1030,8 @@ def _step(pid, args): pstress, ) - f_int = self.nodes.f_int.at[:].set(0) + # f_int = self.nodes.f_int.at[:].set(0) + f_int = self.nodes.f_int mapped_nodes = vmap(self.id_to_node_ids)(particles.element_ids).squeeze(-1) mapped_coords = vmap(self.id_to_node_loc)(particles.element_ids).squeeze(2) mapped_grads = vmap(self.shapefn_grad)( diff --git a/diffmpm/mesh.py b/diffmpm/mesh.py index 15e4fe2..94b861e 100644 --- a/diffmpm/mesh.py +++ b/diffmpm/mesh.py @@ -31,7 +31,7 @@ def __init__(self, config: dict): self.elements: _Element = config["elements"] self.particle_tractions = config["particle_surface_traction"] - # TODO: Convert to using jax directives for loop + # TODO: Change to allow called functions to return outputs def apply_on_elements(self, function: str, args: Tuple = ()): """Apply a given function to elements. diff --git a/diffmpm/node.py b/diffmpm/node.py index a035bd9..eae41ba 100644 --- a/diffmpm/node.py +++ b/diffmpm/node.py @@ -1,9 +1,4 @@ -from typing import Optional, Sized, Tuple - import jax.numpy as jnp -from jax import jit -from jax.tree_util import register_pytree_node_class -from jax.typing import ArrayLike import chex @@ -21,7 +16,7 @@ class _NodesState: f_damp: chex.ArrayDevice -def init_state( +def init_node_state( nnodes: int, loc: chex.ArrayDevice, ): @@ -48,7 +43,7 @@ def init_state( velocity = jnp.zeros_like(loc, dtype=jnp.float32) acceleration = jnp.zeros_like(loc, dtype=jnp.float32) - mass = jnp.ones((loc.shape[0], 1, 1), dtype=jnp.float32) + mass = jnp.zeros((loc.shape[0], 1, 1), dtype=jnp.float32) momentum = jnp.zeros_like(loc, dtype=jnp.float32) f_int = jnp.zeros_like(loc, dtype=jnp.float32) f_ext = jnp.zeros_like(loc, dtype=jnp.float32) @@ -64,3 +59,18 @@ def init_state( f_ext=f_ext, f_damp=f_damp, ) + + +def reset_node_state(state: _NodesState): + mass = state.mass.at[:].set(0) + momentum = state.momentum.at[:].set(0) + f_int = state.f_int.at[:].set(0) + f_ext = state.f_ext.at[:].set(0) + f_damp = state.f_damp.at[:].set(0) + return state.replace( + mass=mass, + momentum=momentum, + f_int=f_int, + f_ext=f_ext, + f_damp=f_damp, + ) diff --git a/diffmpm/scheme.py b/diffmpm/scheme.py index 61a062e..d16f8ca 100644 --- a/diffmpm/scheme.py +++ b/diffmpm/scheme.py @@ -12,6 +12,8 @@ _schemes = ("usf", "usl") +from diffmpm.node import reset_node_state + class _MPMScheme(abc.ABC): def __init__(self, mesh, dt, velocity_update): @@ -21,6 +23,7 @@ def __init__(self, mesh, dt, velocity_update): def compute_nodal_kinematics(self): """Compute nodal kinematics - map mass and momentum to mesh nodes.""" + self.mesh.elements.nodes = reset_node_state(self.mesh.elements.nodes) self.mesh.apply_on_elements("set_particle_element_ids") self.mesh.apply_on_particles("update_natural_coords") self.mesh.apply_on_elements("compute_nodal_mass") diff --git a/tests/test_element.py b/tests/test_element.py index 5f78de5..e465f56 100644 --- a/tests/test_element.py +++ b/tests/test_element.py @@ -109,7 +109,6 @@ def test_element_node_vel(self, elements): def test_compute_nodal_mass(self, elements, particles): particles = particles.replace(mass=particles.mass + 1) - # particles.mass += 1 elements.compute_nodal_mass(particles) true_mass = jnp.ones((4, 1, 1)) assert jnp.all(elements.nodes.mass == true_mass) From f6e18807822b2e2433d2f579b32dd05df29de203 Mon Sep 17 00:00:00 2001 From: Chahak Mehta Date: Sun, 23 Jul 2023 15:45:36 -0700 Subject: [PATCH 14/21] Return changed element arrays, not new state --- diffmpm/constraint.py | 14 ++--- diffmpm/element.py | 133 +++++++++++++++++++++++++++--------------- diffmpm/mesh.py | 25 +++++++- diffmpm/scheme.py | 6 +- tests/test_element.py | 44 +++++++------- 5 files changed, 142 insertions(+), 80 deletions(-) diff --git a/diffmpm/constraint.py b/diffmpm/constraint.py index cb40a3e..a5310cb 100644 --- a/diffmpm/constraint.py +++ b/diffmpm/constraint.py @@ -26,7 +26,7 @@ def tree_unflatten(cls, aux_data, children): del children return cls(*aux_data) - def apply_vel(self, obj, ids): + def apply_vel(self, vel, ids): """Apply constraint values to the passed object. Parameters @@ -37,10 +37,10 @@ def apply_vel(self, obj, ids): The indices of the container `obj` on which the constraint will be applied. """ - velocity = obj.velocity.at[ids, :, self.dir].set(self.velocity) + velocity = vel.at[ids, :, self.dir].set(self.velocity) return velocity - def apply_mom(self, obj, ids): + def apply_mom(self, mom, mass, ids): """Apply constraint values to the passed object. Parameters @@ -51,12 +51,10 @@ def apply_mom(self, obj, ids): The indices of the container `obj` on which the constraint will be applied. """ - momentum = obj.momentum.at[ids, :, self.dir].set( - obj.mass[ids, :, 0] * self.velocity - ) + momentum = mom.at[ids, :, self.dir].set(mass[ids, :, 0] * self.velocity) return momentum - def apply_acc(self, obj, ids): + def apply_acc(self, acc, ids): """Apply constraint values to the passed object. Parameters @@ -67,5 +65,5 @@ def apply_acc(self, obj, ids): The indices of the container `obj` on which the constraint will be applied. """ - acceleration = obj.acceleration.at[ids, :, self.dir].set(0) + acceleration = acc.at[ids, :, self.dir].set(0) return acceleration diff --git a/diffmpm/element.py b/diffmpm/element.py index 3c15ec6..db93871 100644 --- a/diffmpm/element.py +++ b/diffmpm/element.py @@ -164,7 +164,7 @@ def _step(pid, args): ) _, mass, _, _ = lax.fori_loop(0, particles.nparticles, _step, args) # TODO: Return state instead of setting - self.nodes = self.nodes.replace(mass=mass) + return mass, "mass" def compute_nodal_momentum(self, particles: _ParticlesState): r"""Compute the nodal mass based on particle mass. @@ -201,7 +201,7 @@ def _step(pid, args): _, new_momentum, _, _ = lax.fori_loop(0, particles.nparticles, _step, args) new_momentum = jnp.where(jnp.abs(new_momentum) < 1e-12, 0, new_momentum) # TODO: Return state instead of setting - self.nodes = self.nodes.replace(momentum=new_momentum) + return new_momentum, "momentum" def compute_velocity(self, particles: _ParticlesState): """Compute velocity using momentum.""" @@ -216,7 +216,7 @@ def compute_velocity(self, particles: _ParticlesState): velocity, ) # TODO: Return state instead of setting - self.nodes = self.nodes.replace(velocity=velocity) + return velocity, "velocity" def compute_external_force(self, particles: _ParticlesState): r"""Update the nodal external force based on particle f_ext. @@ -252,7 +252,7 @@ def _step(pid, args): ) f_ext, _, _, _ = lax.fori_loop(0, particles.nparticles, _step, args) # TODO: Return state instead of setting - self.nodes = self.nodes.replace(f_ext=f_ext) + return f_ext, "f_ext" def compute_body_force(self, particles: _ParticlesState, gravity: ArrayLike): r"""Update the nodal external force based on particle mass. @@ -289,7 +289,7 @@ def _step(pid, args): ) f_ext, _, _, _, _ = lax.fori_loop(0, particles.nparticles, _step, args) # TODO: Return state instead of setting - self.nodes = self.nodes.replace(f_ext=f_ext) + return f_ext, "f_ext" def apply_concentrated_nodal_forces( self, particles: _ParticlesState, curr_time: float @@ -326,7 +326,7 @@ def _f(x, *, orig): _step_2 = tree_reduce(lambda x, y: x + y, _step_1) f_ext = jnp.where(_step_2 == 0, self.nodes.f_ext, _step_2) # TODO: Return state instead of setting - self.nodes = self.nodes.replace(f_ext=f_ext) + return f_ext, "f_ext" def apply_particle_traction_forces(self, particles: _ParticlesState): """Apply concentrated nodal forces. @@ -348,45 +348,80 @@ def _step(pid, args): args = (self.nodes.f_ext, particles.traction, mapped_positions, mapped_nodes) f_ext, _, _, _ = lax.fori_loop(0, particles.nparticles, _step, args) # TODO: Return state instead of setting - self.nodes = self.nodes.replace(f_ext=f_ext) + return f_ext, "f_ext" - def update_nodal_acceleration_velocity( - self, particles: _ParticlesState, dt: float, *args - ): + def update_nodal_acceleration(self, particles: _ParticlesState, dt: float, *args): """Update the nodal momentum based on total force on nodes.""" total_force = self.nodes.f_int + self.nodes.f_ext + self.nodes.f_damp acceleration = self.nodes.acceleration.at[:].set( jnp.nan_to_num(jnp.divide(total_force, self.nodes.mass)) ) - velocity = self.nodes.velocity.at[:].add(acceleration * dt) - self.nodes = self.nodes.replace(velocity=velocity, acceleration=acceleration) - self.apply_boundary_constraints() - momentum = self.nodes.momentum.at[:].set(self.nodes.mass * velocity) - velocity = jnp.where( - jnp.abs(self.nodes.velocity) < 1e-12, + # velocity = self.nodes.velocity.at[:].add(acceleration * dt) + # self.nodes = self.nodes.replace(velocity=velocity, acceleration=acceleration) + if self.constraints: + acceleration = self._apply_boundary_constraints_acc(acceleration) + # momentum = self.nodes.momentum.at[:].set(self.nodes.mass * velocity) + # velocity = jnp.where( + # jnp.abs(self.nodes.velocity) < 1e-12, + # 0, + # self.nodes.velocity, + # ) + acceleration = jnp.where( + jnp.abs(acceleration) < 1e-12, 0, - self.nodes.velocity, + acceleration, ) - acceleration = jnp.where( - jnp.abs(self.nodes.acceleration) < 1e-12, + # TODO: Return state instead of setting + # self.nodes = self.nodes.replace( + # velocity=velocity, acceleration=acceleration, momentum=momentum + # ) + return acceleration, "acceleration" + + def update_nodal_velocity(self, particles: _ParticlesState, dt: float, *args): + """Update the nodal momentum based on total force on nodes.""" + total_force = self.nodes.f_int + self.nodes.f_ext + self.nodes.f_damp + acceleration = jnp.nan_to_num(jnp.divide(total_force, self.nodes.mass)) + + velocity = self.nodes.velocity + acceleration * dt + if self.constraints: + velocity = self._apply_boundary_constraints_vel(velocity) + velocity = jnp.where( + jnp.abs(velocity) < 1e-12, 0, - self.nodes.acceleration, + velocity, ) + # acceleration = jnp.where( + # jnp.abs(self.nodes.acceleration) < 1e-12, + # 0, + # self.nodes.acceleration, + # ) # TODO: Return state instead of setting - self.nodes = self.nodes.replace( - velocity=velocity, acceleration=acceleration, momentum=momentum + # self.nodes = self.nodes.replace( + # velocity=velocity, acceleration=acceleration, momentum=momentum + # ) + return velocity, "velocity" + + def update_nodal_momentum(self, particles: _ParticlesState, dt: float, *args): + """Update the nodal momentum based on total force on nodes.""" + momentum = self.nodes.momentum.at[:].set(self.nodes.mass * self.nodes.velocity) + momentum = jnp.where( + jnp.abs(momentum) < 1e-12, + 0, + momentum, ) + # TODO: Return state instead of setting + return momentum, "momentum" - def _apply_boundary_constraints_vel(self, *args): + def _apply_boundary_constraints_vel(self, vel, *args): """Apply boundary conditions for nodal velocity.""" # This assumes that the constraints don't have overlapping # conditions. In case it does, only the first constraint will # be applied. - def _func2(constraint, *, nodes): - return constraint[1].apply_vel(nodes, constraint[0]) + def _func2(constraint, *, orig): + return constraint[1].apply_vel(orig, constraint[0]) - partial_func = partial(_func2, nodes=self.nodes) + partial_func = partial(_func2, orig=vel) _out = tree_map( partial_func, self.constraints, is_leaf=lambda x: isinstance(x, tuple) ) @@ -400,26 +435,26 @@ def _func2(constraint, *, nodes): def _f(x, *, orig): return jnp.where(x == orig, jnp.nan, x) - _pf = partial(_f, orig=self.nodes.velocity) + _pf = partial(_f, orig=vel) _step_1 = tree_map(_pf, _out) vel = tree_reduce( lambda x, y: jnp.where(jnp.isnan(y), x, y), - [self.nodes.velocity, _step_1], + [vel, _step_1], ) # TODO: Return state instead of setting - self.nodes = self.nodes.replace(velocity=vel) + return vel - def _apply_boundary_constraints_mom(self, *args): + def _apply_boundary_constraints_mom(self, mom, mass, *args): """Apply boundary conditions for nodal momentum.""" # This assumes that the constraints don't have overlapping # conditions. In case it does, only the first constraint will # be applied. - def _func2(constraint, *, nodes): - return constraint[1].apply_mom(nodes, constraint[0]) + def _func2(constraint, *, mom, mass): + return constraint[1].apply_mom(mom, mass, constraint[0]) - partial_func = partial(_func2, nodes=self.nodes) + partial_func = partial(_func2, mom=mom, mass=mass) _out = tree_map( partial_func, self.constraints, is_leaf=lambda x: isinstance(x, tuple) ) @@ -433,26 +468,26 @@ def _func2(constraint, *, nodes): def _f(x, *, orig): return jnp.where(x == orig, jnp.nan, x) - _pf = partial(_f, orig=self.nodes.momentum) + _pf = partial(_f, orig=mom) _step_1 = tree_map(_pf, _out) mom = tree_reduce( lambda x, y: jnp.where(jnp.isnan(y), x, y), - [self.nodes.momentum, _step_1], + [mom, _step_1], ) # TODO: Return state instead of setting - self.nodes = self.nodes.replace(momentum=mom) + return mom - def _apply_boundary_constraints_acc(self, *args): + def _apply_boundary_constraints_acc(self, orig, *args): """Apply boundary conditions for nodal acceleration.""" # This assumes that the constraints don't have overlapping # conditions. In case it does, only the first constraint will # be applied. - def _func2(constraint, *, nodes): - return constraint[1].apply_acc(nodes, constraint[0]) + def _func2(constraint, *, orig): + return constraint[1].apply_acc(orig, constraint[0]) - partial_func = partial(_func2, nodes=self.nodes) + partial_func = partial(_func2, orig=orig) _out = tree_map( partial_func, self.constraints, is_leaf=lambda x: isinstance(x, tuple) ) @@ -466,21 +501,25 @@ def _func2(constraint, *, nodes): def _f(x, *, orig): return jnp.where(x == orig, jnp.nan, x) - _pf = partial(_f, orig=self.nodes.acceleration) + _pf = partial(_f, orig=orig) _step_1 = tree_map(_pf, _out) acc = tree_reduce( lambda x, y: jnp.where(jnp.isnan(y), x, y), - [self.nodes.acceleration, _step_1], + [orig, _step_1], ) # TODO: Return state instead of setting - self.nodes = self.nodes.replace(acceleration=acc) + return acc def apply_boundary_constraints(self, *args): if self.constraints: - self._apply_boundary_constraints_vel(*args) - self._apply_boundary_constraints_mom(*args) - self._apply_boundary_constraints_acc(*args) + vel = self._apply_boundary_constraints_vel(self.nodes.velocity, *args) + mom = self._apply_boundary_constraints_mom( + self.nodes.momentum, self.nodes.mass, *args + ) + acc = self._apply_boundary_constraints_acc(self.nodes.acceleration, *args) + + return self.nodes.replace(velocity=vel, momentum=mom, acceleration=acc) @register_pytree_node_class @@ -1047,7 +1086,7 @@ def _step(pid, args): ) f_int, _, _, _, _ = lax.fori_loop(0, particles.nparticles, _step, args) # TODO: Return state instead of setting - self.nodes = self.nodes.replace(f_int=f_int) + return f_int, "f_int" def compute_volume(self, *args): """Compute volume of all elements.""" diff --git a/diffmpm/mesh.py b/diffmpm/mesh.py index 94b861e..2c06004 100644 --- a/diffmpm/mesh.py +++ b/diffmpm/mesh.py @@ -7,6 +7,7 @@ from jax.tree_util import register_pytree_node_class, tree_map from diffmpm.element import _Element +from diffmpm.node import _NodesState from diffmpm.particle import _ParticlesState import diffmpm.particle as dpart from diffmpm.forces import ParticleTraction @@ -55,8 +56,30 @@ def _func(particles, *, func, fargs): ) if function == "set_particle_element_ids": self.particles = _out + elif function == "apply_boundary_constraints": + self.elements.nodes = _out[0] + elif _out[0] is not None: + _temp = tree_util.tree_transpose( + tree_util.tree_structure([0 for e in _out]), + tree_util.tree_structure(_out[0]), + _out, + ) + + def reduce_attr(state_1, state_2, *, orig): + new_val = state_1 + state_2 - orig + return new_val + + attr = _temp[1][0] + p_reduce_attr = partial( + reduce_attr, + attr=attr, + orig=getattr(self.elements.nodes, attr), + ) + new_val = tree_util.tree_reduce( + p_reduce_attr, _temp[0], is_leaf=lambda x: isinstance(x, _NodesState) + ) + self.elements.nodes = self.elements.nodes.replace(**{attr: new_val}) - # TODO: Convert to using jax directives for loop def apply_on_particles(self, function: str, args: Tuple = ()): """Apply a given function to particles. diff --git a/diffmpm/scheme.py b/diffmpm/scheme.py index d16f8ca..385dd7b 100644 --- a/diffmpm/scheme.py +++ b/diffmpm/scheme.py @@ -60,9 +60,9 @@ def compute_forces(self, gravity: ArrayLike, step: int): def compute_particle_kinematics(self): """Compute particle location, acceleration and velocity.""" - self.mesh.apply_on_elements( - "update_nodal_acceleration_velocity", args=(self.dt,) - ) + self.mesh.apply_on_elements("update_nodal_acceleration", args=(self.dt,)) + self.mesh.apply_on_elements("update_nodal_velocity", args=(self.dt,)) + self.mesh.apply_on_elements("update_nodal_momentum", args=(self.dt,)) self.mesh.apply_on_particles( "update_position_velocity", args=(self.dt, self.velocity_update), diff --git a/tests/test_element.py b/tests/test_element.py index e465f56..72adc6c 100644 --- a/tests/test_element.py +++ b/tests/test_element.py @@ -109,21 +109,21 @@ def test_element_node_vel(self, elements): def test_compute_nodal_mass(self, elements, particles): particles = particles.replace(mass=particles.mass + 1) - elements.compute_nodal_mass(particles) + nodal_mass, _ = elements.compute_nodal_mass(particles) true_mass = jnp.ones((4, 1, 1)) - assert jnp.all(elements.nodes.mass == true_mass) + assert jnp.all(nodal_mass == true_mass) def test_compute_nodal_momentum(self, elements, particles): particles = particles.replace(velocity=particles.velocity + 1) - elements.compute_nodal_momentum(particles) + nodal_momentum, _ = elements.compute_nodal_momentum(particles) true_momentum = jnp.ones((4, 1, 1)) * 0.5 - assert jnp.all(elements.nodes.momentum == true_momentum) + assert jnp.all(nodal_momentum == true_momentum) def test_compute_external_force(self, elements, particles): particles = particles.replace(f_ext=particles.f_ext + 1) - elements.compute_external_force(particles) + nodal_f_ext, _ = elements.compute_external_force(particles) true_fext = jnp.ones((4, 1, 1)) * 0.5 - assert jnp.all(elements.nodes.f_ext == true_fext) + assert jnp.all(nodal_f_ext == true_fext) @pytest.mark.parametrize( "gravity, expected", @@ -137,8 +137,8 @@ def test_compute_external_force(self, elements, particles): ) def test_compute_body_force(self, elements, particles, gravity, expected): particles = particles.replace(mass=particles.mass + 1) - elements.compute_body_force(particles, gravity) - assert jnp.all(elements.nodes.f_ext == expected) + nodal_f_ext, _ = elements.compute_body_force(particles, gravity) + assert jnp.all(nodal_f_ext == expected) def test_apply_concentrated_nodal_force(self, particles): cnf_1 = NodalForce( @@ -157,10 +157,9 @@ def test_apply_concentrated_nodal_force(self, particles): (1, 1), 1, 1, [], concentrated_nodal_forces=[cnf_1, cnf_2] ) elements.nodes = elements.nodes.replace(f_ext=elements.nodes.f_ext + 2) - elements.apply_concentrated_nodal_forces(particles, 1) + nodal_f_ext, _ = elements.apply_concentrated_nodal_forces(particles, 1) assert jnp.all( - elements.nodes.f_ext - == jnp.array([[3, 2], [2, 2], [3, 3], [2, 2]]).reshape(4, 1, 2) + nodal_f_ext == jnp.array([[3, 2], [2, 2], [3, 3], [2, 2]]).reshape(4, 1, 2) ) def test_apply_boundary_constraints(self): @@ -170,9 +169,9 @@ def test_apply_boundary_constraints(self): ] elements = Quadrilateral4Node((1, 1), 1, (1.0, 1.0), cons) elements.nodes = elements.nodes.replace(velocity=elements.nodes.velocity + 1) - elements.apply_boundary_constraints() + node_state = elements.apply_boundary_constraints() assert jnp.all( - elements.nodes.velocity + node_state.velocity == jnp.array([[0, 2], [0, 1], [1, 1], [1, 1]]).reshape(4, 1, 2) ) @@ -180,17 +179,20 @@ def test_update_nodal_acceleration_velocity(self, elements, particles): f_ext = elements.nodes.f_ext + jnp.array([1, 0]) mass = elements.nodes.mass.at[:].set(2) elements.nodes = elements.nodes.replace(mass=mass, f_ext=f_ext) - elements.update_nodal_acceleration_velocity(particles, 0.1) + nodal_acc, _ = elements.update_nodal_acceleration(particles, 0.1) assert jnp.allclose( - elements.nodes.acceleration, + nodal_acc, jnp.array([[0.5, 0.0], [0.5, 0], [0.5, 0], [0.5, 0]]), ) + nodal_vel, _ = elements.update_nodal_velocity(particles, 0.1) assert jnp.allclose( - elements.nodes.velocity, + nodal_vel, jnp.array([[0.05, 0.0], [0.05, 0], [0.05, 0], [0.05, 0]]), ) + elements.nodes = elements.nodes.replace(velocity=nodal_vel) + nodal_mom, _ = elements.update_nodal_momentum(particles, 0.1) assert jnp.allclose( - elements.nodes.momentum, + nodal_mom, jnp.array([[0.1, 0.0], [0.1, 0], [0.1, 0], [0.1, 0]]), ) @@ -202,9 +204,9 @@ def test_set_particle_element_ids(self, elements, particles): def test_compute_internal_force(self, elements, particles): particles = dpar.compute_volume(particles, elements, 1) particles = particles.replace(stress=particles.stress + 1) - elements.compute_internal_force(particles) + nodal_f_int, _ = elements.compute_internal_force(particles) assert jnp.allclose( - elements.nodes.f_int, + nodal_f_int, jnp.array([[1, 1], [0, 0], [0, 0], [-1, -1]]).reshape(4, 1, 2), ) @@ -214,8 +216,8 @@ def test_compute_volume(self, elements): def test_apply_particle_traction_forces(self, elements, particles): particles = particles.replace(traction=particles.traction + jnp.array([1, 0])) - elements.apply_particle_traction_forces(particles) + nodal_f_ext, _ = elements.apply_particle_traction_forces(particles) assert jnp.allclose( - elements.nodes.f_ext, + nodal_f_ext, jnp.array([[0.5, 0], [0.5, 0], [0.5, 0], [0.5, 0]]).reshape(4, 1, 2), ) From 5f3b5f73b21f025466fdef20cda4f602a9b15759 Mon Sep 17 00:00:00 2001 From: Chahak Mehta Date: Sun, 30 Jul 2023 23:14:49 -0700 Subject: [PATCH 15/21] Working on completely stateless design --- .../mpm-nodal-forces.toml | 2 +- .../uniaxial_nodal_forces/test_benchmark.py | 4 + .../mpm-particle-traction.toml | 2 +- .../uniaxial_stress/mpm-uniaxial-stress.toml | 2 +- .../2d/uniaxial_stress/test_benchmark.py | 4 + diffmpm/element.py | 1834 +++++++++-------- diffmpm/explicit.py | 447 ++++ diffmpm/io.py | 7 +- diffmpm/mesh.py | 23 +- diffmpm/node.py | 20 +- diffmpm/particle.py | 205 +- pyproject.toml | 7 +- tests/test_particle.py | 4 +- 13 files changed, 1635 insertions(+), 926 deletions(-) create mode 100644 diffmpm/explicit.py diff --git a/benchmarks/2d/uniaxial_nodal_forces/mpm-nodal-forces.toml b/benchmarks/2d/uniaxial_nodal_forces/mpm-nodal-forces.toml index 93a9386..ffa4a96 100644 --- a/benchmarks/2d/uniaxial_nodal_forces/mpm-nodal-forces.toml +++ b/benchmarks/2d/uniaxial_nodal_forces/mpm-nodal-forces.toml @@ -33,7 +33,7 @@ type = "generator" nelements = [3, 1] element_length = [0.1, 0.1] particle_element_ids = [0] -element = "Quadrilateral4Node" +element = "Quad4N" entity_sets = "entity_sets.json" [[mesh.constraints]] diff --git a/benchmarks/2d/uniaxial_nodal_forces/test_benchmark.py b/benchmarks/2d/uniaxial_nodal_forces/test_benchmark.py index 5038150..2ece328 100644 --- a/benchmarks/2d/uniaxial_nodal_forces/test_benchmark.py +++ b/benchmarks/2d/uniaxial_nodal_forces/test_benchmark.py @@ -35,3 +35,7 @@ def test_benchmarks(): result = jnp.load("results/uniaxial-nodal-forces/particles_0990.npz") assert jnp.round(result["stress"][0, :, 0].min() - 0.9999990078443788, 5) == 0.0 assert jnp.round(result["stress"][0, :, 0].max() - 0.9999990292713694, 5) == 0.0 + + +if __name__ == "__main__": + test_benchmarks() diff --git a/benchmarks/2d/uniaxial_particle_traction/mpm-particle-traction.toml b/benchmarks/2d/uniaxial_particle_traction/mpm-particle-traction.toml index 067cbb9..4bfec27 100644 --- a/benchmarks/2d/uniaxial_particle_traction/mpm-particle-traction.toml +++ b/benchmarks/2d/uniaxial_particle_traction/mpm-particle-traction.toml @@ -33,7 +33,7 @@ type = "generator" nelements = [3, 1] element_length = [0.1, 0.1] particle_element_ids = [0] -element = "Quadrilateral4Node" +element = "Quad4N" entity_sets = "entity_sets.json" [[mesh.constraints]] diff --git a/benchmarks/2d/uniaxial_stress/mpm-uniaxial-stress.toml b/benchmarks/2d/uniaxial_stress/mpm-uniaxial-stress.toml index bf15148..39c141c 100644 --- a/benchmarks/2d/uniaxial_stress/mpm-uniaxial-stress.toml +++ b/benchmarks/2d/uniaxial_stress/mpm-uniaxial-stress.toml @@ -29,7 +29,7 @@ type = "generator" nelements = [1, 1] element_length = [1, 1] particle_element_ids = [0] -element = "Quadrilateral4Node" +element = "Quad4N" entity_sets = "entity_sets.json" [[mesh.constraints]] diff --git a/benchmarks/2d/uniaxial_stress/test_benchmark.py b/benchmarks/2d/uniaxial_stress/test_benchmark.py index ee589dd..f17e4fe 100644 --- a/benchmarks/2d/uniaxial_stress/test_benchmark.py +++ b/benchmarks/2d/uniaxial_stress/test_benchmark.py @@ -23,3 +23,7 @@ def test_benchmarks(): assert jnp.round(result["stress"][0, :, 1].max() - true_stress_yy, 8) == 0.0 assert jnp.round(result["stress"][0, :, 0].max() - true_stress_xx, 8) == 0.0 + + +if __name__ == "__main__": + test_benchmarks() diff --git a/diffmpm/element.py b/diffmpm/element.py index db93871..a068c57 100644 --- a/diffmpm/element.py +++ b/diffmpm/element.py @@ -10,7 +10,7 @@ import jax.numpy as jnp from jax import Array, jacobian, jit, lax, tree_util, vmap -from jax.tree_util import register_pytree_node_class, tree_map, tree_reduce +from jax.tree_util import register_pytree_node_class, tree_map, tree_reduce, Partial from jax.typing import ArrayLike from diffmpm.constraint import Constraint @@ -18,44 +18,92 @@ from diffmpm.node import _NodesState, init_node_state import chex -__all__ = ["_Element", "Linear1D", "Quadrilateral4Node"] - @chex.dataclass() -class _ElementState: +class _ElementsState: nodes: _NodesState total_elements: int - concentrated_nodal_forces: Sequence volume: chex.ArrayDevice + constraints: Sequence[Tuple[ArrayLike, Constraint]] + concentrated_nodal_forces: Sequence[NodalForce] -class _Element(abc.ABC): - """Base element class that is inherited by all types of Elements.""" +@chex.dataclass() +class Quad4NState(_ElementsState): + nelements: chex.ArrayDevice + el_len: chex.ArrayDevice - nodes: _NodesState - total_elements: int - concentrated_nodal_forces: Sequence - volume: Array - @abc.abstractmethod - def id_to_node_ids(self, id: ArrayLike) -> Array: - """Node IDs corresponding to element `id`. +@chex.dataclass() +class Quad4N: + total_elements: int - This method is implemented by each of the subclass. + def init_state( + self, + nelements: int, + total_elements: int, + el_len: float, + constraints: Sequence[Tuple[ArrayLike, Constraint]], + nodes: Optional[_NodesState] = None, + concentrated_nodal_forces: Sequence = [], + initialized: Optional[bool] = None, + volume: Optional[ArrayLike] = None, + ) -> Quad4NState: + """Initialize Linear1D. Parameters ---------- - id : int - Element ID. - - Returns - ------- - ArrayLike - Nodal IDs of the element. + nelements : int + Number of elements. + total_elements : int + Total number of elements (product of all elements of `nelements`) + el_len : float + Length of each element. + constraints: list + A list of constraints where each element is a tuple of + type `(node_ids, diffmpm.Constraint)`. Here, `node_ids` + correspond to the node IDs where `diffmpm.Constraint` + should be applied. + nodes : Nodes, Optional + Nodes in the element object. + concentrated_nodal_forces: list + A list of `diffmpm.forces.NodalForce`s that are to be + applied. + initialized: bool, None + `True` if the class has been initialized, `None` if not. + This is required like this for using JAX flattening. + volume: ArrayLike + Volume of the elements. """ - ... + nelements = jnp.asarray(nelements) + el_len = jnp.asarray(el_len) + + total_nodes = jnp.prod(nelements + 1) + coords = jnp.asarray( + list( + itertools.product( + jnp.arange(nelements[1] + 1), + jnp.arange(nelements[0] + 1), + ) + ) + ) + node_locations = (jnp.asarray([coords[:, 1], coords[:, 0]]).T * el_len).reshape( + -1, 1, 2 + ) + nodes = init_node_state(int(total_nodes), node_locations) + + volume = jnp.ones((total_elements, 1, 1)) + return Quad4NState( + nodes=nodes, + total_elements=total_elements, + concentrated_nodal_forces=concentrated_nodal_forces, + volume=volume, + constraints=constraints, + nelements=nelements, + el_len=el_len, + ) - def id_to_node_loc(self, id: ArrayLike) -> Array: + def id_to_node_loc(self, elements: _ElementState, id: ArrayLike) -> Array: """Node locations corresponding to element `id`. Parameters @@ -69,10 +117,10 @@ def id_to_node_loc(self, id: ArrayLike) -> Array: Nodal locations for the element. Shape of returned array is `(nodes_in_element, 1, ndim)` """ - node_ids = self.id_to_node_ids(id).squeeze() - return self.nodes.loc[node_ids] + node_ids = self.id_to_node_ids(elements.nelements[0], id).squeeze() + return elements.nodes.loc[node_ids] - def id_to_node_vel(self, id: ArrayLike) -> Array: + def id_to_node_vel(self, elements: _ElementState, id: ArrayLike) -> Array: """Node velocities corresponding to element `id`. Parameters @@ -86,649 +134,356 @@ def id_to_node_vel(self, id: ArrayLike) -> Array: Nodal velocities for the element. Shape of returned array is `(nodes_in_element, 1, ndim)` """ - node_ids = self.id_to_node_ids(id).squeeze() - return self.nodes.velocity[node_ids] - - def tree_flatten(self): - children = (self.nodes, self.volume) - aux_data = ( - self.nelements, - self.total_elements, - self.el_len, - self.constraints, - self.concentrated_nodal_forces, - self.initialized, - ) - return children, aux_data - - @classmethod - def tree_unflatten(cls, aux_data, children): - return cls( - aux_data[0], - aux_data[1], - aux_data[2], - aux_data[3], - nodes=children[0], - concentrated_nodal_forces=aux_data[4], - initialized=aux_data[5], - volume=children[1], - ) - - @abc.abstractmethod - def shapefn(self, xi: ArrayLike): - """Evaluate Shape function for element type.""" - ... - - @abc.abstractmethod - def shapefn_grad(self, xi: ArrayLike, coords: ArrayLike): - """Evaluate gradient of shape function for element type.""" - ... - - @abc.abstractmethod - def set_particle_element_ids(self, particles: _ParticlesState): - """Set the element IDs that particles are present in.""" - ... + node_ids = self.id_to_node_ids(elements.nelements[0], id).squeeze() + return elements.nodes.velocity[node_ids] - # Mapping from particles to nodes (P2G) - def compute_nodal_mass(self, particles: _ParticlesState): - r"""Compute the nodal mass based on particle mass. + def id_to_node_ids(self, nelements_x, id: ArrayLike): + """Node IDs corresponding to element `id`. - The nodal mass is updated as a sum of particle mass for - all particles mapped to the node. + 3----2 + | | + 0----1 - \[ - (m)_i = \sum_p N_i(x_p) m_p - \] + Node ids are returned in the order as shown in the figure. Parameters ---------- - particles: diffmpm.particle.Particles - Particles to map to the nodal values. - """ - - @jit - def _step(pid, args): - pmass, mass, mapped_pos, el_nodes = args - mass = mass.at[el_nodes[pid]].add(pmass[pid] * mapped_pos[pid]) - return pmass, mass, mapped_pos, el_nodes + id : int + Element ID. - # mass = self.nodes.mass.at[:].set(0) - mass = self.nodes.mass - mapped_positions = self.shapefn(particles.reference_loc) - mapped_nodes = vmap(jit(self.id_to_node_ids))(particles.element_ids).squeeze(-1) - args = ( - particles.mass, - mass, - mapped_positions, - mapped_nodes, + Returns + ------- + ArrayLike + Nodal IDs of the element. Shape of returned + array is (4, 1) + """ + lower_left = (id // nelements_x) * (nelements_x + 1) + id % nelements_x + result = jnp.asarray( + [ + lower_left, + lower_left + 1, + lower_left + nelements_x + 2, + lower_left + nelements_x + 1, + ] ) - _, mass, _, _ = lax.fori_loop(0, particles.nparticles, _step, args) - # TODO: Return state instead of setting - return mass, "mass" + return result.reshape(4, 1) - def compute_nodal_momentum(self, particles: _ParticlesState): - r"""Compute the nodal mass based on particle mass. + @classmethod + def _get_mapped_nodes(cls, id, nelements_x): + """Node IDs corresponding to element `id`. - The nodal mass is updated as a sum of particle mass for - all particles mapped to the node. + 3----2 + | | + 0----1 - \[ - (mv)_i = \sum_p N_i(x_p) (mv)_p - \] + Node ids are returned in the order as shown in the figure. Parameters ---------- - particles: diffmpm.particle.Particles - Particles to map to the nodal values. - """ - - @jit - def _step(pid, args): - pmom, mom, mapped_pos, el_nodes = args - new_mom = mom.at[el_nodes[pid]].add(mapped_pos[pid] @ pmom[pid]) - return pmom, new_mom, mapped_pos, el_nodes + id : int + Element ID. - # curr_mom = self.nodes.momentum.at[:].set(0) - curr_mom = self.nodes.momentum - mapped_positions = self.shapefn(particles.reference_loc) - mapped_nodes = vmap(jit(self.id_to_node_ids))(particles.element_ids).squeeze(-1) - args = ( - particles.mass * particles.velocity, - curr_mom, - mapped_positions, - mapped_nodes, + Returns + ------- + ArrayLike + Nodal IDs of the element. Shape of returned + array is (4, 1) + """ + lower_left = (id // nelements_x) * (nelements_x + 1) + id % nelements_x + result = jnp.asarray( + [ + lower_left, + lower_left + 1, + lower_left + nelements_x + 2, + lower_left + nelements_x + 1, + ] ) - _, new_momentum, _, _ = lax.fori_loop(0, particles.nparticles, _step, args) - new_momentum = jnp.where(jnp.abs(new_momentum) < 1e-12, 0, new_momentum) - # TODO: Return state instead of setting - return new_momentum, "momentum" + return result.reshape(4, 1) - def compute_velocity(self, particles: _ParticlesState): - """Compute velocity using momentum.""" - velocity = jnp.where( - self.nodes.mass == 0, - self.nodes.velocity, - self.nodes.momentum / self.nodes.mass, - ) - velocity = jnp.where( - jnp.abs(velocity) < 1e-12, - 0, - velocity, - ) - # TODO: Return state instead of setting - return velocity, "velocity" + def shapefn(self, xi: ArrayLike): + """Evaluate linear shape function. - def compute_external_force(self, particles: _ParticlesState): - r"""Update the nodal external force based on particle f_ext. + Parameters + ---------- + xi : float, array_like + Locations of particles in natural coordinates to evaluate + the function at. Expected shape is (npoints, 1, ndim) - The nodal force is updated as a sum of particle external - force for all particles mapped to the node. + Returns + ------- + array_like + Evaluated shape function values. The shape of the returned + array will depend on the input shape. For example, in the linear + case, if the input is a scalar, the returned array will be of + the shape `(1, 4, 1)` but if the input is a vector then the output will + be of the shape `(len(x), 4, 1)`. + """ + xi = jnp.asarray(xi) + if xi.ndim != 3: + raise ValueError( + f"`xi` should be of size (npoints, 1, ndim); found {xi.shape}" + ) + result = jnp.array( + [ + 0.25 * (1 - xi[:, :, 0]) * (1 - xi[:, :, 1]), + 0.25 * (1 + xi[:, :, 0]) * (1 - xi[:, :, 1]), + 0.25 * (1 + xi[:, :, 0]) * (1 + xi[:, :, 1]), + 0.25 * (1 - xi[:, :, 0]) * (1 + xi[:, :, 1]), + ] + ) + result = result.transpose(1, 0, 2)[..., jnp.newaxis] + return result - \[ - f_{ext})_i = \sum_p N_i(x_p) f_{ext} - \] + @classmethod + def _shapefn(cls, xi: ArrayLike): + """Evaluate linear shape function. Parameters ---------- - particles: diffmpm.particle.Particles - Particles to map to the nodal values. - """ - - @jit - def _step(pid, args): - f_ext, pf_ext, mapped_pos, el_nodes = args - f_ext = f_ext.at[el_nodes[pid]].add(mapped_pos[pid] @ pf_ext[pid]) - return f_ext, pf_ext, mapped_pos, el_nodes + xi : float, array_like + Locations of particles in natural coordinates to evaluate + the function at. Expected shape is (npoints, 1, ndim) - # f_ext = self.nodes.f_ext.at[:].set(0) - f_ext = self.nodes.f_ext - mapped_positions = self.shapefn(particles.reference_loc) - mapped_nodes = vmap(jit(self.id_to_node_ids))(particles.element_ids).squeeze(-1) - args = ( - f_ext, - particles.f_ext, - mapped_positions, - mapped_nodes, + Returns + ------- + array_like + Evaluated shape function values. The shape of the returned + array will depend on the input shape. For example, in the linear + case, if the input is a scalar, the returned array will be of + the shape `(1, 4, 1)` but if the input is a vector then the output will + be of the shape `(len(x), 4, 1)`. + """ + xi = jnp.asarray(xi) + if xi.ndim != 3: + raise ValueError( + f"`xi` should be of size (npoints, 1, ndim); found {xi.shape}" + ) + result = jnp.array( + [ + 0.25 * (1 - xi[:, :, 0]) * (1 - xi[:, :, 1]), + 0.25 * (1 + xi[:, :, 0]) * (1 - xi[:, :, 1]), + 0.25 * (1 + xi[:, :, 0]) * (1 + xi[:, :, 1]), + 0.25 * (1 - xi[:, :, 0]) * (1 + xi[:, :, 1]), + ] ) - f_ext, _, _, _ = lax.fori_loop(0, particles.nparticles, _step, args) - # TODO: Return state instead of setting - return f_ext, "f_ext" - - def compute_body_force(self, particles: _ParticlesState, gravity: ArrayLike): - r"""Update the nodal external force based on particle mass. + result = result.transpose(1, 0, 2)[..., jnp.newaxis] + return result - The nodal force is updated as a sum of particle body - force for all particles mapped to th + @classmethod + def _shapefn_natural_grad(cls, xi: ArrayLike): + """Calculate the gradient of shape function. - \[ - (f_{ext})_i = (f_{ext})_i + \sum_p N_i(x_p) m_p g - \] + This calculation is done in the natural coordinates. Parameters ---------- - particles: diffmpm.particle._ParticlesState - Particles to map to the nodal values. - """ - - @jit - def _step(pid, args): - f_ext, pmass, mapped_pos, el_nodes, gravity = args - f_ext = f_ext.at[el_nodes[pid]].add( - mapped_pos[pid] @ (pmass[pid] * gravity) - ) - return f_ext, pmass, mapped_pos, el_nodes, gravity + x : float, array_like + Locations of particles in natural coordinates to evaluate + the function at. - mapped_positions = self.shapefn(particles.reference_loc) - mapped_nodes = vmap(jit(self.id_to_node_ids))(particles.element_ids).squeeze(-1) - args = ( - self.nodes.f_ext, - particles.mass, - mapped_positions, - mapped_nodes, - gravity, + Returns + ------- + array_like + Evaluated gradient values of the shape function. The shape of + the returned array will depend on the input shape. For example, + in the linear case, if the input is a scalar, the returned array + will be of the shape `(4, 2)`. + """ + # result = vmap(jacobian(self.shapefn))(xi[..., jnp.newaxis]).squeeze() + xi = jnp.asarray(xi) + xi = xi.squeeze() + result = jnp.array( + [ + [-0.25 * (1 - xi[1]), -0.25 * (1 - xi[0])], + [0.25 * (1 - xi[1]), -0.25 * (1 + xi[0])], + [0.25 * (1 + xi[1]), 0.25 * (1 + xi[0])], + [-0.25 * (1 + xi[1]), 0.25 * (1 - xi[0])], + ], ) - f_ext, _, _, _, _ = lax.fori_loop(0, particles.nparticles, _step, args) - # TODO: Return state instead of setting - return f_ext, "f_ext" + return result - def apply_concentrated_nodal_forces( - self, particles: _ParticlesState, curr_time: float - ): - """Apply concentrated nodal forces. + def shapefn_grad(self, xi: ArrayLike, coords: ArrayLike): + """Gradient of shape function in physical coordinates. Parameters ---------- - particles: _ParticlesState - Particles in the simulation. - curr_time: float - Current time in the simulation. - """ - - def _func(cnf, *, f_ext): - factor = cnf.function.value(curr_time) - f_ext = f_ext.at[cnf.node_ids, 0, cnf.dir].add(factor * cnf.force) - return f_ext + xi : float, array_like + Locations of particles to evaluate in natural coordinates. + Expected shape `(npoints, 1, ndim)`. + coords : array_like + Nodal coordinates to transform by. Expected shape + `(npoints, 1, ndim)` - if self.concentrated_nodal_forces: - partial_func = partial(_func, f_ext=self.nodes.f_ext) - _out = tree_map( - partial_func, - self.concentrated_nodal_forces, - is_leaf=lambda x: isinstance(x, NodalForce), + Returns + ------- + array_like + Gradient of the shape function in physical coordinates at `xi` + """ + xi = jnp.asarray(xi) + coords = jnp.asarray(coords) + if xi.ndim != 3: + raise ValueError( + f"`x` should be of size (npoints, 1, ndim); found {xi.shape}" ) + grad_sf = self._shapefn_natural_grad(xi) + _jacobian = grad_sf.T @ coords.squeeze() - def _f(x, *, orig): - return jnp.where(x == orig, 0, x) - - # This assumes that the nodal forces are not overlapping, i.e. - # no node will be acted by 2 forces in the same direction. - _step_1 = tree_map(partial(_f, orig=self.nodes.f_ext), _out) - _step_2 = tree_reduce(lambda x, y: x + y, _step_1) - f_ext = jnp.where(_step_2 == 0, self.nodes.f_ext, _step_2) - # TODO: Return state instead of setting - return f_ext, "f_ext" + result = grad_sf @ jnp.linalg.inv(_jacobian).T + return result - def apply_particle_traction_forces(self, particles: _ParticlesState): - """Apply concentrated nodal forces. + @classmethod + def _shapefn_grad(cls, xi: ArrayLike, coords: ArrayLike): + """Gradient of shape function in physical coordinates. Parameters ---------- - particles: Particles - Particles in the simulation. + xi : float, array_like + Locations of particles to evaluate in natural coordinates. + Expected shape `(npoints, 1, ndim)`. + coords : array_like + Nodal coordinates to transform by. Expected shape + `(npoints, 1, ndim)` + + Returns + ------- + array_like + Gradient of the shape function in physical coordinates at `xi` """ + xi = jnp.asarray(xi) + coords = jnp.asarray(coords) + if xi.ndim != 3: + raise ValueError( + f"`x` should be of size (npoints, 1, ndim); found {xi.shape}" + ) + grad_sf = cls._shapefn_natural_grad(xi) + _jacobian = grad_sf.T @ coords.squeeze() - @jit - def _step(pid, args): - f_ext, ptraction, mapped_pos, el_nodes = args - f_ext = f_ext.at[el_nodes[pid]].add(mapped_pos[pid] @ ptraction[pid]) - return f_ext, ptraction, mapped_pos, el_nodes + result = grad_sf @ jnp.linalg.inv(_jacobian).T + return result - mapped_positions = self.shapefn(particles.reference_loc) - mapped_nodes = vmap(jit(self.id_to_node_ids))(particles.element_ids).squeeze(-1) - args = (self.nodes.f_ext, particles.traction, mapped_positions, mapped_nodes) - f_ext, _, _, _ = lax.fori_loop(0, particles.nparticles, _step, args) - # TODO: Return state instead of setting - return f_ext, "f_ext" + @classmethod + def _get_particles_element_ids(cls, particles, elements): + """Set the element IDs for the particles. - def update_nodal_acceleration(self, particles: _ParticlesState, dt: float, *args): - """Update the nodal momentum based on total force on nodes.""" - total_force = self.nodes.f_int + self.nodes.f_ext + self.nodes.f_damp - acceleration = self.nodes.acceleration.at[:].set( - jnp.nan_to_num(jnp.divide(total_force, self.nodes.mass)) - ) - # velocity = self.nodes.velocity.at[:].add(acceleration * dt) - # self.nodes = self.nodes.replace(velocity=velocity, acceleration=acceleration) - if self.constraints: - acceleration = self._apply_boundary_constraints_acc(acceleration) - # momentum = self.nodes.momentum.at[:].set(self.nodes.mass * velocity) - # velocity = jnp.where( - # jnp.abs(self.nodes.velocity) < 1e-12, - # 0, - # self.nodes.velocity, - # ) - acceleration = jnp.where( - jnp.abs(acceleration) < 1e-12, - 0, - acceleration, - ) - # TODO: Return state instead of setting - # self.nodes = self.nodes.replace( - # velocity=velocity, acceleration=acceleration, momentum=momentum - # ) - return acceleration, "acceleration" + If the particle doesn't lie between the boundaries of any + element, it sets the element index to -1. + """ - def update_nodal_velocity(self, particles: _ParticlesState, dt: float, *args): - """Update the nodal momentum based on total force on nodes.""" - total_force = self.nodes.f_int + self.nodes.f_ext + self.nodes.f_damp - acceleration = jnp.nan_to_num(jnp.divide(total_force, self.nodes.mass)) + def f(x, *, loc, nelements): + xidl = (loc[:, :, 0] <= x[0, 0]).nonzero(size=loc.shape[0], fill_value=-1)[ + 0 + ] + yidl = (loc[:, :, 1] <= x[0, 1]).nonzero(size=loc.shape[0], fill_value=-1)[ + 0 + ] + lower_left = jnp.where(jnp.isin(xidl, yidl), xidl, -1).max() + element_id = lower_left - lower_left // (nelements + 1) + return element_id - velocity = self.nodes.velocity + acceleration * dt - if self.constraints: - velocity = self._apply_boundary_constraints_vel(velocity) - velocity = jnp.where( - jnp.abs(velocity) < 1e-12, - 0, - velocity, - ) - # acceleration = jnp.where( - # jnp.abs(self.nodes.acceleration) < 1e-12, - # 0, - # self.nodes.acceleration, - # ) - # TODO: Return state instead of setting - # self.nodes = self.nodes.replace( - # velocity=velocity, acceleration=acceleration, momentum=momentum - # ) - return velocity, "velocity" + pf = partial(f, loc=elements.nodes.loc, nelements=elements.nelements[0]) + ids = vmap(pf)(particles.loc) + return ids - def update_nodal_momentum(self, particles: _ParticlesState, dt: float, *args): - """Update the nodal momentum based on total force on nodes.""" - momentum = self.nodes.momentum.at[:].set(self.nodes.mass * self.nodes.velocity) - momentum = jnp.where( - jnp.abs(momentum) < 1e-12, - 0, - momentum, - ) - # TODO: Return state instead of setting - return momentum, "momentum" + def set_particle_element_ids( + self, elements: _ElementsState, particles: _ParticlesState + ): + """Set the element IDs for the particles. - def _apply_boundary_constraints_vel(self, vel, *args): - """Apply boundary conditions for nodal velocity.""" + If the particle doesn't lie between the boundaries of any + element, it sets the element index to -1. + """ - # This assumes that the constraints don't have overlapping - # conditions. In case it does, only the first constraint will - # be applied. - def _func2(constraint, *, orig): - return constraint[1].apply_vel(orig, constraint[0]) + @jit + def f(x, *, loc, nelements): + xidl = (loc[:, :, 0] <= x[0, 0]).nonzero(size=loc.shape[0], fill_value=-1)[ + 0 + ] + yidl = (loc[:, :, 1] <= x[0, 1]).nonzero(size=loc.shape[0], fill_value=-1)[ + 0 + ] + lower_left = jnp.where(jnp.isin(xidl, yidl), xidl, -1).max() + element_id = lower_left - lower_left // (nelements + 1) + return element_id - partial_func = partial(_func2, orig=vel) - _out = tree_map( - partial_func, self.constraints, is_leaf=lambda x: isinstance(x, tuple) - ) + pf = partial(f, loc=elements.nodes.loc, nelements=elements.nelements[0]) + ids = vmap(pf)(particles.loc) + return particles.replace(element_ids=ids) - # _temp = tree_util.tree_transpose( - # tree_util.tree_structure([0 for e in _out]), - # tree_util.tree_structure(_out[0]), - # _out, - # ) + def compute_internal_force( + self, elements: _ElementState, particles: _ParticlesState + ): + r"""Update the nodal internal force based on particle mass. - def _f(x, *, orig): - return jnp.where(x == orig, jnp.nan, x) + The nodal force is updated as a sum of internal forces for + all particles mapped to the node. - _pf = partial(_f, orig=vel) - _step_1 = tree_map(_pf, _out) - vel = tree_reduce( - lambda x, y: jnp.where(jnp.isnan(y), x, y), - [vel, _step_1], - ) + \[ + (f_{int})_i = -\sum_p V_p \sigma_p \nabla N_i(x_p) + \] - # TODO: Return state instead of setting - return vel + where \(\sigma_p\) is the stress at particle \(p\). - def _apply_boundary_constraints_mom(self, mom, mass, *args): - """Apply boundary conditions for nodal momentum.""" + Parameters + ---------- + particles: diffmpm.particle._ParticlesState + Particles to map to the nodal values. + """ - # This assumes that the constraints don't have overlapping - # conditions. In case it does, only the first constraint will - # be applied. - def _func2(constraint, *, mom, mass): - return constraint[1].apply_mom(mom, mass, constraint[0]) + @jit + def _step(pid, args): + ( + f_int, + pvol, + mapped_grads, + el_nodes, + pstress, + ) = args + force = jnp.zeros((mapped_grads.shape[1], 1, 2)) + force = force.at[:, 0, 0].set( + mapped_grads[pid][:, 0] * pstress[pid][0] + + mapped_grads[pid][:, 1] * pstress[pid][3] + ) + force = force.at[:, 0, 1].set( + mapped_grads[pid][:, 1] * pstress[pid][1] + + mapped_grads[pid][:, 0] * pstress[pid][3] + ) + update = -pvol[pid] * force + f_int = f_int.at[el_nodes[pid]].add(update) + return ( + f_int, + pvol, + mapped_grads, + el_nodes, + pstress, + ) - partial_func = partial(_func2, mom=mom, mass=mass) - _out = tree_map( - partial_func, self.constraints, is_leaf=lambda x: isinstance(x, tuple) + # f_int = self.nodes.f_int.at[:].set(0) + f_int = elements.nodes.f_int + mapped_nodes = vmap(Partial(self.id_to_node_ids, elements.nelements[0]))( + particles.element_ids + ).squeeze(-1) + mapped_coords = vmap(Partial(self.id_to_node_loc, elements))( + particles.element_ids + ).squeeze(2) + mapped_grads = vmap(jit(self.shapefn_grad))( + particles.reference_loc[:, jnp.newaxis, ...], + mapped_coords, ) - - # _temp = tree_util.tree_transpose( - # tree_util.tree_structure([0 for e in _out]), - # tree_util.tree_structure(_out[0]), - # _out, - # ) - - def _f(x, *, orig): - return jnp.where(x == orig, jnp.nan, x) - - _pf = partial(_f, orig=mom) - _step_1 = tree_map(_pf, _out) - mom = tree_reduce( - lambda x, y: jnp.where(jnp.isnan(y), x, y), - [mom, _step_1], + args = ( + f_int, + particles.volume, + mapped_grads, + mapped_nodes, + particles.stress, ) + f_int, _, _, _, _ = lax.fori_loop(0, particles.nparticles, _step, args) + return f_int, "f_int" - # TODO: Return state instead of setting - return mom - - def _apply_boundary_constraints_acc(self, orig, *args): - """Apply boundary conditions for nodal acceleration.""" - - # This assumes that the constraints don't have overlapping - # conditions. In case it does, only the first constraint will - # be applied. - def _func2(constraint, *, orig): - return constraint[1].apply_acc(orig, constraint[0]) - - partial_func = partial(_func2, orig=orig) - _out = tree_map( - partial_func, self.constraints, is_leaf=lambda x: isinstance(x, tuple) - ) - - # _temp = tree_util.tree_transpose( - # tree_util.tree_structure([0 for e in _out]), - # tree_util.tree_structure(_out[0]), - # _out, - # ) - - def _f(x, *, orig): - return jnp.where(x == orig, jnp.nan, x) - - _pf = partial(_f, orig=orig) - _step_1 = tree_map(_pf, _out) - acc = tree_reduce( - lambda x, y: jnp.where(jnp.isnan(y), x, y), - [orig, _step_1], - ) - - # TODO: Return state instead of setting - return acc - - def apply_boundary_constraints(self, *args): - if self.constraints: - vel = self._apply_boundary_constraints_vel(self.nodes.velocity, *args) - mom = self._apply_boundary_constraints_mom( - self.nodes.momentum, self.nodes.mass, *args - ) - acc = self._apply_boundary_constraints_acc(self.nodes.acceleration, *args) - - return self.nodes.replace(velocity=vel, momentum=mom, acceleration=acc) - - -@register_pytree_node_class -class Linear1D(_Element): - """Container for 1D line elements (and nodes). - - Element ID: 0 1 2 3 - Mesh: +-----+-----+-----+-----+ - Node IDs: 0 1 2 3 4 - - where - - + : Nodes - +-----+ : An element - - """ - - def __init__( - self, - nelements: int, - total_elements: int, - el_len: float, - constraints: Sequence[Tuple[ArrayLike, Constraint]], - nodes: Optional[Nodes] = None, - concentrated_nodal_forces: Sequence = [], - initialized: Optional[bool] = None, - volume: Optional[ArrayLike] = None, - ): - """Initialize Linear1D. - - Parameters - ---------- - nelements : int - Number of elements. - total_elements : int - Total number of elements (same as `nelements` for 1D) - el_len : float - Length of each element. - constraints: list - A list of constraints where each element is a tuple of type - `(node_ids, diffmpm.Constraint)`. Here, `node_ids` correspond to - the node IDs where `diffmpm.Constraint` should be applied. - nodes : Nodes, Optional - Nodes in the element object. - concentrated_nodal_forces: list - A list of `diffmpm.forces.NodalForce`s that are to be - applied. - initialized: bool, None - `True` if the class has been initialized, `None` if not. - This is required like this for using JAX flattening. - volume: ArrayLike - Volume of the elements. - """ - self.nelements = nelements - self.total_elements = nelements - self.el_len = el_len - if nodes is None: - self.nodes = Nodes( - nelements + 1, - jnp.arange(nelements + 1).reshape(-1, 1, 1) * el_len, - ) - else: - self.nodes = nodes - - # self.boundary_nodes = boundary_nodes - self.constraints = constraints - self.concentrated_nodal_forces = concentrated_nodal_forces - if initialized is None: - self.volume = jnp.ones((self.total_elements, 1, 1)) - else: - self.volume = jnp.asarray(volume) - self.initialized = True - - def id_to_node_ids(self, id: ArrayLike): - """Node IDs corresponding to element `id`. - - Parameters - ---------- - id : int - Element ID. - - Returns - ------- - ArrayLike - Nodal IDs of the element. Shape of returned - array is `(2, 1)` - """ - return jnp.array([id, id + 1]).reshape(2, 1) - - def shapefn(self, xi: ArrayLike): - """Evaluate linear shape function. - - Parameters - ---------- - xi : float, array_like - Locations of particles in natural coordinates to evaluate - the function at. Expected shape is `(npoints, 1, ndim)` - - Returns - ------- - array_like - Evaluated shape function values. The shape of the returned - array will depend on the input shape. For example, in the linear - case, if the input is a scalar, the returned array will be of - the shape `(1, 2, 1)` but if the input is a vector then the output will - be of the shape `(len(x), 2, 1)`. - """ - xi = jnp.asarray(xi) - if xi.ndim != 3: - raise ValueError( - f"`xi` should be of size (npoints, 1, ndim); found {xi.shape}" - ) - result = jnp.array([0.5 * (1 - xi), 0.5 * (1 + xi)]).transpose(1, 0, 2, 3) - return result - - def _shapefn_natural_grad(self, xi: ArrayLike): - """Calculate the gradient of shape function. - - This calculation is done in the natural coordinates. - - Parameters - ---------- - x : float, array_like - Locations of particles in natural coordinates to evaluate - the function at. - - Returns - ------- - array_like - Evaluated gradient values of the shape function. The shape of - the returned array will depend on the input shape. For example, - in the linear case, if the input is a scalar, the returned array - will be of the shape `(2, 1)`. - """ - xi = jnp.asarray(xi) - result = vmap(jacobian(self.shapefn))(xi[..., jnp.newaxis]).squeeze() - - # TODO: The following code tries to evaluate vmap even if - # the predicate condition is true, not sure why. - # result = lax.cond( - # jnp.isscalar(x), - # jacobian(self.shapefn), - # vmap(jacobian(self.shapefn)), - # xi - # ) - return result.reshape(2, 1) - - def shapefn_grad(self, xi: ArrayLike, coords: ArrayLike): - """Gradient of shape function in physical coordinates. - - Parameters - ---------- - xi : float, array_like - Locations of particles to evaluate in natural coordinates. - Expected shape `(npoints, 1, ndim)`. - coords : array_like - Nodal coordinates to transform by. Expected shape - `(npoints, 1, ndim)` - - Returns - ------- - array_like - Gradient of the shape function in physical coordinates at `xi` - """ - xi = jnp.asarray(xi) - coords = jnp.asarray(coords) - if xi.ndim != 3: - raise ValueError( - f"`x` should be of size (npoints, 1, ndim); found {xi.shape}" - ) - grad_sf = self._shapefn_natural_grad(xi) - _jacobian = grad_sf.T @ coords - - result = grad_sf @ jnp.linalg.inv(_jacobian).T - return result - - def set_particle_element_ids(self, particles): - """Set the element IDs for the particles. - - If the particle doesn't lie between the boundaries of any - element, it sets the element index to -1. - """ - - @jit - def f(x): - idl = ( - len(self.nodes.loc) - - 1 - - jnp.asarray(self.nodes.loc[::-1] <= x).nonzero(size=1, fill_value=-1)[ - 0 - ][-1] - ) - idg = ( - jnp.asarray(self.nodes.loc > x).nonzero(size=1, fill_value=-1)[0][0] - 1 - ) - return (idl, idg) - - ids = vmap(f)(particles.loc) - particles.element_ids = jnp.where( - ids[0] == ids[1], ids[0], jnp.ones_like(ids[0]) * -1 - ) - - def compute_volume(self, *args): - """Compute volume of all elements.""" - vol = jnp.ediff1d(self.nodes.loc) - self.volume = jnp.ones((self.total_elements, 1, 1)) * vol - - def compute_internal_force(self, particles): - r"""Update the nodal internal force based on particle mass. + @classmethod + def _compute_internal_force( + cls, nf_int, nloc, mapped_node_ids, pxi, pvol, pstress, nparticles + ): + r"""Update the nodal internal force based on particle mass. The nodal force is updated as a sum of internal forces for all particles mapped to the node. @@ -745,6 +500,7 @@ def compute_internal_force(self, particles): Particles to map to the nodal values. """ + @jit def _step(pid, args): ( f_int, @@ -753,10 +509,17 @@ def _step(pid, args): el_nodes, pstress, ) = args - # TODO: correct matrix multiplication for n-d - # update = -(pvol[pid]) * pstress[pid] @ mapped_grads[pid] - update = -pvol[pid] * pstress[pid][0] * mapped_grads[pid] - f_int = f_int.at[el_nodes[pid]].add(update[..., jnp.newaxis]) + force = jnp.zeros((mapped_grads.shape[1], 1, 2)) + force = force.at[:, 0, 0].set( + mapped_grads[pid][:, 0] * pstress[pid][0] + + mapped_grads[pid][:, 1] * pstress[pid][3] + ) + force = force.at[:, 0, 1].set( + mapped_grads[pid][:, 1] * pstress[pid][1] + + mapped_grads[pid][:, 0] * pstress[pid][3] + ) + update = -pvol[pid] * force + f_int = f_int.at[el_nodes[pid]].add(update) return ( f_int, pvol, @@ -765,276 +528,291 @@ def _step(pid, args): pstress, ) - self.nodes.f_int = self.nodes.f_int.at[:].set(0) - mapped_nodes = vmap(jit(self.id_to_node_ids))(particles.element_ids).squeeze(-1) - mapped_coords = vmap(jit(self.id_to_node_loc))(particles.element_ids).squeeze(2) - mapped_grads = vmap(jit(self.shapefn_grad))( - particles.reference_loc[:, jnp.newaxis, ...], + # f_int = self.nodes.f_int.at[:].set(0) + # f_int = elements.nodes.f_int + mapped_nodes = mapped_node_ids.squeeze(-1) + mapped_coords = nloc[mapped_nodes].squeeze(2) + mapped_grads = vmap(jit(cls._shapefn_grad))( + pxi[:, jnp.newaxis, ...], mapped_coords, ) args = ( - self.nodes.f_int, - particles.volume, + nf_int, + pvol, mapped_grads, mapped_nodes, - particles.stress, + pstress, ) - self.nodes.f_int, _, _, _, _ = lax.fori_loop( - 0, particles.nparticles, _step, args - ) - + f_int, _, _, _, _ = lax.fori_loop(0, nparticles, _step, args) + return f_int -@register_pytree_node_class -class Quadrilateral4Node(_Element): - r"""Container for 2D quadrilateral elements with 4 nodes. - - Nodes and elements are numbered as - - 15 +---+---+---+---+ 19 - | 8 | 9 | 10| 11| - 10 +---+---+---+---+ 14 - | 4 | 5 | 6 | 7 | - 5 +---+---+---+---+ 9 - | 0 | 1 | 2 | 3 | - +---+---+---+---+ - 0 1 2 3 4 - - where + # Mapping from particles to nodes (P2G) + @classmethod + def _compute_nodal_mass(cls, mass, pmass, pxi, peids, mapped_node_ids, nparticles): + r"""Compute the nodal mass based on particle mass. - + : Nodes - +---+ - | | : An element - +---+ - """ + The nodal mass is updated as a sum of particle mass for + all particles mapped to the node. - def __init__( - self, - nelements: int, - total_elements: int, - el_len: float, - constraints: Sequence[Tuple[ArrayLike, Constraint]], - nodes: Optional[_NodesState] = None, - concentrated_nodal_forces: Sequence = [], - initialized: Optional[bool] = None, - volume: Optional[ArrayLike] = None, - ) -> None: - """Initialize Linear1D. + \[ + (m)_i = \sum_p N_i(x_p) m_p + \] Parameters ---------- - nelements : int - Number of elements. - total_elements : int - Total number of elements (product of all elements of `nelements`) - el_len : float - Length of each element. - constraints: list - A list of constraints where each element is a tuple of - type `(node_ids, diffmpm.Constraint)`. Here, `node_ids` - correspond to the node IDs where `diffmpm.Constraint` - should be applied. - nodes : Nodes, Optional - Nodes in the element object. - concentrated_nodal_forces: list - A list of `diffmpm.forces.NodalForce`s that are to be - applied. - initialized: bool, None - `True` if the class has been initialized, `None` if not. - This is required like this for using JAX flattening. - volume: ArrayLike - Volume of the elements. + particles: diffmpm.particle.Particles + Particles to map to the nodal values. """ - self.nelements = jnp.asarray(nelements) - self.el_len = jnp.asarray(el_len) - self.total_elements = total_elements - - if nodes is None: - total_nodes = jnp.prod(self.nelements + 1) - coords = jnp.asarray( - list( - itertools.product( - jnp.arange(self.nelements[1] + 1), - jnp.arange(self.nelements[0] + 1), - ) - ) - ) - node_locations = ( - jnp.asarray([coords[:, 1], coords[:, 0]]).T * self.el_len - ).reshape(-1, 1, 2) - self.nodes = init_node_state(int(total_nodes), node_locations) - else: - self.nodes = nodes - - self.constraints = constraints - self.concentrated_nodal_forces = concentrated_nodal_forces - if initialized is None: - self.volume = jnp.ones((self.total_elements, 1, 1)) - else: - self.volume = jnp.asarray(volume) - self.initialized = True - - def id_to_node_ids(self, id: ArrayLike): - """Node IDs corresponding to element `id`. - 3----2 - | | - 0----1 + @jit + def _step(pid, args): + pmass, mass, mapped_pos, el_nodes = args + mass = mass.at[el_nodes[pid]].add(pmass[pid] * mapped_pos[pid]) + return pmass, mass, mapped_pos, el_nodes - Node ids are returned in the order as shown in the figure. + mapped_positions = cls._shapefn(pxi) + mapped_nodes = mapped_node_ids.squeeze(-1) + args = ( + pmass, + mass, + mapped_positions, + mapped_nodes, + ) + _, mass, _, _ = lax.fori_loop(0, nparticles, _step, args) + return mass + + def compute_nodal_mass(self, elements, particles: _ParticlesState): + r"""Compute the nodal mass based on particle mass. + + The nodal mass is updated as a sum of particle mass for + all particles mapped to the node. + + \[ + (m)_i = \sum_p N_i(x_p) m_p + \] Parameters ---------- - id : int - Element ID. - - Returns - ------- - ArrayLike - Nodal IDs of the element. Shape of returned - array is (4, 1) + particles: diffmpm.particle.Particles + Particles to map to the nodal values. """ - lower_left = (id // self.nelements[0]) * ( - self.nelements[0] + 1 - ) + id % self.nelements[0] - result = jnp.asarray( - [ - lower_left, - lower_left + 1, - lower_left + self.nelements[0] + 2, - lower_left + self.nelements[0] + 1, - ] + + @jit + def _step(pid, args): + pmass, mass, mapped_pos, el_nodes = args + mass = mass.at[el_nodes[pid]].add(pmass[pid] * mapped_pos[pid]) + return pmass, mass, mapped_pos, el_nodes + + # mass = self.nodes.mass.at[:].set(0) + mass = elements.nodes.mass + mapped_positions = self.shapefn(particles.reference_loc) + mapped_nodes = vmap(Partial(self.id_to_node_ids, elements.nelements[0]))( + particles.element_ids + ).squeeze(-1) + args = ( + particles.mass, + mass, + mapped_positions, + mapped_nodes, ) - return result.reshape(4, 1) + _, mass, _, _ = lax.fori_loop(0, particles.nparticles, _step, args) + # TODO: Return state instead of setting + return mass, "mass" - def shapefn(self, xi: ArrayLike): - """Evaluate linear shape function. + @classmethod + def _compute_nodal_momentum( + cls, nmom, pmass, pvel, pxi, peids, mapped_node_ids, nparticles + ): + r"""Compute the nodal mass based on particle mass. + + The nodal mass is updated as a sum of particle mass for + all particles mapped to the node. + + \[ + (mv)_i = \sum_p N_i(x_p) (mv)_p + \] Parameters ---------- - xi : float, array_like - Locations of particles in natural coordinates to evaluate - the function at. Expected shape is (npoints, 1, ndim) - - Returns - ------- - array_like - Evaluated shape function values. The shape of the returned - array will depend on the input shape. For example, in the linear - case, if the input is a scalar, the returned array will be of - the shape `(1, 4, 1)` but if the input is a vector then the output will - be of the shape `(len(x), 4, 1)`. + particles: diffmpm.particle.Particles + Particles to map to the nodal values. """ - xi = jnp.asarray(xi) - if xi.ndim != 3: - raise ValueError( - f"`xi` should be of size (npoints, 1, ndim); found {xi.shape}" - ) - result = jnp.array( - [ - 0.25 * (1 - xi[:, :, 0]) * (1 - xi[:, :, 1]), - 0.25 * (1 + xi[:, :, 0]) * (1 - xi[:, :, 1]), - 0.25 * (1 + xi[:, :, 0]) * (1 + xi[:, :, 1]), - 0.25 * (1 - xi[:, :, 0]) * (1 + xi[:, :, 1]), - ] + + @jit + def _step(pid, args): + pmom, mom, mapped_pos, el_nodes = args + new_mom = mom.at[el_nodes[pid]].add(mapped_pos[pid] @ pmom[pid]) + return pmom, new_mom, mapped_pos, el_nodes + + # curr_mom = elements.nodes.momentum.at[:].set(0) + # curr_mom = elements.nodes.momentum + mapped_nodes = mapped_node_ids.squeeze(-1) + mapped_positions = cls._shapefn(pxi) + args = ( + pmass * pvel, + nmom, + mapped_positions, + mapped_nodes, ) - result = result.transpose(1, 0, 2)[..., jnp.newaxis] - return result + _, new_momentum, _, _ = lax.fori_loop(0, nparticles, _step, args) + new_momentum = jnp.where(jnp.abs(new_momentum) < 1e-12, 0, new_momentum) + return new_momentum - def _shapefn_natural_grad(self, xi: ArrayLike): - """Calculate the gradient of shape function. + def compute_nodal_momentum(self, elements, particles: _ParticlesState): + r"""Compute the nodal mass based on particle mass. - This calculation is done in the natural coordinates. + The nodal mass is updated as a sum of particle mass for + all particles mapped to the node. + + \[ + (mv)_i = \sum_p N_i(x_p) (mv)_p + \] Parameters ---------- - x : float, array_like - Locations of particles in natural coordinates to evaluate - the function at. - - Returns - ------- - array_like - Evaluated gradient values of the shape function. The shape of - the returned array will depend on the input shape. For example, - in the linear case, if the input is a scalar, the returned array - will be of the shape `(4, 2)`. + particles: diffmpm.particle.Particles + Particles to map to the nodal values. """ - # result = vmap(jacobian(self.shapefn))(xi[..., jnp.newaxis]).squeeze() - xi = jnp.asarray(xi) - xi = xi.squeeze() - result = jnp.array( - [ - [-0.25 * (1 - xi[1]), -0.25 * (1 - xi[0])], - [0.25 * (1 - xi[1]), -0.25 * (1 + xi[0])], - [0.25 * (1 + xi[1]), 0.25 * (1 + xi[0])], - [-0.25 * (1 + xi[1]), 0.25 * (1 - xi[0])], - ], + + @jit + def _step(pid, args): + pmom, mom, mapped_pos, el_nodes = args + new_mom = mom.at[el_nodes[pid]].add(mapped_pos[pid] @ pmom[pid]) + return pmom, new_mom, mapped_pos, el_nodes + + # curr_mom = elements.nodes.momentum.at[:].set(0) + curr_mom = elements.nodes.momentum + mapped_positions = self.shapefn(particles.reference_loc) + mapped_nodes = vmap(Partial(self.id_to_node_ids, elements.nelements[0]))( + particles.element_ids + ).squeeze(-1) + args = ( + particles.mass * particles.velocity, + curr_mom, + mapped_positions, + mapped_nodes, ) - return result + _, new_momentum, _, _ = lax.fori_loop(0, particles.nparticles, _step, args) + new_momentum = jnp.where(jnp.abs(new_momentum) < 1e-12, 0, new_momentum) + # TODO: Return state instead of setting + return new_momentum, "momentum" - def shapefn_grad(self, xi: ArrayLike, coords: ArrayLike): - """Gradient of shape function in physical coordinates. + @classmethod + def _compute_nodal_velocity(cls, nmass, nmom, nvel): + """Compute velocity using momentum.""" + velocity = jnp.where( + nmass == 0, + nvel, + nmom / nmass, + ) + velocity = jnp.where( + jnp.abs(velocity) < 1e-12, + 0, + velocity, + ) + # TODO: Return state instead of setting + return velocity + + def compute_velocity(self, elements, particles: _ParticlesState): + """Compute velocity using momentum.""" + velocity = jnp.where( + elements.nodes.mass == 0, + elements.nodes.velocity, + elements.nodes.momentum / elements.nodes.mass, + ) + velocity = jnp.where( + jnp.abs(velocity) < 1e-12, + 0, + velocity, + ) + # TODO: Return state instead of setting + return velocity, "velocity" + + def compute_external_force(self, elements, particles: _ParticlesState): + r"""Update the nodal external force based on particle f_ext. + + The nodal force is updated as a sum of particle external + force for all particles mapped to the node. + + \[ + f_{ext})_i = \sum_p N_i(x_p) f_{ext} + \] Parameters ---------- - xi : float, array_like - Locations of particles to evaluate in natural coordinates. - Expected shape `(npoints, 1, ndim)`. - coords : array_like - Nodal coordinates to transform by. Expected shape - `(npoints, 1, ndim)` - - Returns - ------- - array_like - Gradient of the shape function in physical coordinates at `xi` + particles: diffmpm.particle.Particles + Particles to map to the nodal values. """ - xi = jnp.asarray(xi) - coords = jnp.asarray(coords) - if xi.ndim != 3: - raise ValueError( - f"`x` should be of size (npoints, 1, ndim); found {xi.shape}" - ) - grad_sf = self._shapefn_natural_grad(xi) - _jacobian = grad_sf.T @ coords.squeeze() - result = grad_sf @ jnp.linalg.inv(_jacobian).T - return result + @jit + def _step(pid, args): + f_ext, pf_ext, mapped_pos, el_nodes = args + f_ext = f_ext.at[el_nodes[pid]].add(mapped_pos[pid] @ pf_ext[pid]) + return f_ext, pf_ext, mapped_pos, el_nodes - def set_particle_element_ids(self, particles: _ParticlesState): - """Set the element IDs for the particles. + # f_ext = elements.nodes.f_ext.at[:].set(0) + f_ext = elements.nodes.f_ext + mapped_positions = self.shapefn(particles.reference_loc) + mapped_nodes = vmap(Partial(self.id_to_node_ids, elements.nelements[0]))( + particles.element_ids + ).squeeze(-1) + args = ( + f_ext, + particles.f_ext, + mapped_positions, + mapped_nodes, + ) + f_ext, _, _, _ = lax.fori_loop(0, particles.nparticles, _step, args) + # TODO: Return state instead of setting + return f_ext, "f_ext" - If the particle doesn't lie between the boundaries of any - element, it sets the element index to -1. + @classmethod + def _compute_external_force(cls, f_ext, pf_ext, pxi, nparticles, mapped_node_ids): + r"""Update the nodal external force based on particle f_ext. + + The nodal force is updated as a sum of particle external + force for all particles mapped to the node. + + \[ + f_{ext})_i = \sum_p N_i(x_p) f_{ext} + \] + + Parameters + ---------- + particles: diffmpm.particle.Particles + Particles to map to the nodal values. """ @jit - def f(x): - xidl = (self.nodes.loc[:, :, 0] <= x[0, 0]).nonzero( - size=len(self.nodes.loc), fill_value=-1 - )[0] - yidl = (self.nodes.loc[:, :, 1] <= x[0, 1]).nonzero( - size=len(self.nodes.loc), fill_value=-1 - )[0] - lower_left = jnp.where(jnp.isin(xidl, yidl), xidl, -1).max() - element_id = lower_left - lower_left // (self.nelements[0] + 1) - return element_id + def _step(pid, args): + f_ext, pf_ext, mapped_pos, el_nodes = args + f_ext = f_ext.at[el_nodes[pid]].add(mapped_pos[pid] @ pf_ext[pid]) + return f_ext, pf_ext, mapped_pos, el_nodes - ids = vmap(f)(particles.loc) - return particles.replace(element_ids=ids) + # f_ext = elements.nodes.f_ext.at[:].set(0) + mapped_positions = cls._shapefn(pxi) + mapped_nodes = mapped_node_ids.squeeze(-1) + args = ( + f_ext, + pf_ext, + mapped_positions, + mapped_nodes, + ) + f_ext, _, _, _ = lax.fori_loop(0, nparticles, _step, args) + return f_ext - def compute_internal_force(self, particles: _ParticlesState): - r"""Update the nodal internal force based on particle mass. + def compute_body_force( + self, elements, particles: _ParticlesState, gravity: ArrayLike + ): + r"""Update the nodal external force based on particle mass. - The nodal force is updated as a sum of internal forces for - all particles mapped to the node. + The nodal force is updated as a sum of particle body + force for all particles mapped to th \[ - (f_{int})_i = -\sum_p V_p \sigma_p \nabla N_i(x_p) + (f_{ext})_i = (f_{ext})_i + \sum_p N_i(x_p) m_p g \] - where \(\sigma_p\) is the stress at particle \(p\). - Parameters ---------- particles: diffmpm.particle._ParticlesState @@ -1043,55 +821,373 @@ def compute_internal_force(self, particles: _ParticlesState): @jit def _step(pid, args): - ( - f_int, - pvol, - mapped_grads, - el_nodes, - pstress, - ) = args - force = jnp.zeros((mapped_grads.shape[1], 1, 2)) - force = force.at[:, 0, 0].set( - mapped_grads[pid][:, 0] * pstress[pid][0] - + mapped_grads[pid][:, 1] * pstress[pid][3] + f_ext, pmass, mapped_pos, el_nodes, gravity = args + f_ext = f_ext.at[el_nodes[pid]].add( + mapped_pos[pid] @ (pmass[pid] * gravity) ) - force = force.at[:, 0, 1].set( - mapped_grads[pid][:, 1] * pstress[pid][1] - + mapped_grads[pid][:, 0] * pstress[pid][3] + return f_ext, pmass, mapped_pos, el_nodes, gravity + + mapped_positions = self.shapefn(particles.reference_loc) + mapped_nodes = vmap(Partial(self.id_to_node_ids, elements.nelements[0]))( + particles.element_ids + ).squeeze(-1) + args = ( + elements.nodes.f_ext, + particles.mass, + mapped_positions, + mapped_nodes, + gravity, + ) + f_ext, _, _, _, _ = lax.fori_loop(0, particles.nparticles, _step, args) + # TODO: Return state instead of setting + return f_ext, "f_ext" + + @classmethod + def _compute_body_force( + cls, nf_ext, pmass, pxi, mapped_node_ids, nparticles, gravity: ArrayLike + ): + r"""Update the nodal external force based on particle mass. + + The nodal force is updated as a sum of particle body + force for all particles mapped to th + + \[ + (f_{ext})_i = (f_{ext})_i + \sum_p N_i(x_p) m_p g + \] + + Parameters + ---------- + particles: diffmpm.particle._ParticlesState + Particles to map to the nodal values. + """ + + @jit + def _step(pid, args): + f_ext, pmass, mapped_pos, el_nodes, gravity = args + f_ext = f_ext.at[el_nodes[pid]].add( + mapped_pos[pid] @ (pmass[pid] * gravity) ) - update = -pvol[pid] * force - f_int = f_int.at[el_nodes[pid]].add(update) - return ( - f_int, - pvol, - mapped_grads, - el_nodes, - pstress, + return f_ext, pmass, mapped_pos, el_nodes, gravity + + mapped_positions = cls._shapefn(pxi) + mapped_nodes = mapped_node_ids.squeeze(-1) + args = ( + nf_ext, + pmass, + mapped_positions, + mapped_nodes, + gravity, + ) + f_ext, _, _, _, _ = lax.fori_loop(0, nparticles, _step, args) + return f_ext + + def apply_concentrated_nodal_forces( + self, elements, particles: _ParticlesState, curr_time: float + ): + """Apply concentrated nodal forces. + + Parameters + ---------- + particles: _ParticlesState + Particles in the simulation. + curr_time: float + Current time in the simulation. + """ + + def _func(cnf, *, f_ext): + factor = cnf.function.value(curr_time) + f_ext = f_ext.at[cnf.node_ids, 0, cnf.dir].add(factor * cnf.force) + return f_ext + + if elements.concentrated_nodal_forces: + partial_func = partial(_func, f_ext=elements.nodes.f_ext) + _out = tree_map( + partial_func, + elements.concentrated_nodal_forces, + is_leaf=lambda x: isinstance(x, NodalForce), ) - # f_int = self.nodes.f_int.at[:].set(0) - f_int = self.nodes.f_int - mapped_nodes = vmap(self.id_to_node_ids)(particles.element_ids).squeeze(-1) - mapped_coords = vmap(self.id_to_node_loc)(particles.element_ids).squeeze(2) - mapped_grads = vmap(self.shapefn_grad)( - particles.reference_loc[:, jnp.newaxis, ...], - mapped_coords, + def _f(x, *, orig): + return jnp.where(x == orig, 0, x) + + # This assumes that the nodal forces are not overlapping, i.e. + # no node will be acted by 2 forces in the same direction. + _step_1 = tree_map(partial(_f, orig=elements.nodes.f_ext), _out) + _step_2 = tree_reduce(lambda x, y: x + y, _step_1) + f_ext = jnp.where(_step_2 == 0, elements.nodes.f_ext, _step_2) + # TODO: Return state instead of setting + return f_ext, "f_ext" + + @classmethod + def _apply_concentrated_nodal_forces( + self, nf_ext, concentrated_forces, curr_time: float + ): + """Apply concentrated nodal forces. + + Parameters + ---------- + particles: _ParticlesState + Particles in the simulation. + curr_time: float + Current time in the simulation. + """ + + def _func(cnf, f_ext): + factor = cnf.function.value(curr_time) + f_ext = f_ext.at[cnf.node_ids, 0, cnf.dir].add(factor * cnf.force) + return f_ext + + _out = tree_map( + _func, + concentrated_forces, + [nf_ext] * len(concentrated_forces), + is_leaf=lambda x: isinstance(x, NodalForce) or isinstance(x, Array), ) + + def _f(x, *, orig): + return jnp.where(x == orig, 0, x) + + # This assumes that the nodal forces are not overlapping, i.e. + # no node will be acted by 2 forces in the same direction. + _step_1 = tree_map(partial(_f, orig=nf_ext), _out) + _step_2 = tree_reduce(lambda x, y: x + y, _step_1) + f_ext = jnp.where(_step_2 == 0, nf_ext, _step_2) + return f_ext + + def apply_particle_traction_forces(self, elements, particles: _ParticlesState): + """Apply concentrated nodal forces. + + Parameters + ---------- + particles: Particles + Particles in the simulation. + """ + + @jit + def _step(pid, args): + f_ext, ptraction, mapped_pos, el_nodes = args + f_ext = f_ext.at[el_nodes[pid]].add(mapped_pos[pid] @ ptraction[pid]) + return f_ext, ptraction, mapped_pos, el_nodes + + mapped_positions = self.shapefn(particles.reference_loc) + mapped_nodes = vmap(Partial(self.id_to_node_ids, elements.nelements[0]))( + particles.element_ids + ).squeeze(-1) args = ( - f_int, - particles.volume, - mapped_grads, + elements.nodes.f_ext, + particles.traction, + mapped_positions, mapped_nodes, - particles.stress, ) - f_int, _, _, _, _ = lax.fori_loop(0, particles.nparticles, _step, args) + f_ext, _, _, _ = lax.fori_loop(0, particles.nparticles, _step, args) # TODO: Return state instead of setting - return f_int, "f_int" + return f_ext, "f_ext" + + def update_nodal_acceleration( + self, elements, particles: _ParticlesState, dt: float, *args + ): + """Update the nodal momentum based on total force on nodes.""" + total_force = ( + elements.nodes.f_int + elements.nodes.f_ext + elements.nodes.f_damp + ) + acceleration = elements.nodes.acceleration.at[:].set( + jnp.nan_to_num(jnp.divide(total_force, elements.nodes.mass)) + ) + if elements.constraints: + acceleration = self._apply_boundary_constraints_acc(elements, acceleration) + acceleration = jnp.where( + jnp.abs(acceleration) < 1e-12, + 0, + acceleration, + ) + return acceleration, "acceleration" + + @classmethod + def _update_nodal_acceleration( + cls, + total_force, + nacc, + nmass, + constraints, + tol, + ): + """Update the nodal momentum based on total force on nodes.""" + acceleration = jnp.nan_to_num(jnp.divide(total_force, nmass)) + if constraints: + acceleration = cls._apply_boundary_constraints_acc( + constraints, acceleration + ) + acceleration = jnp.where( + jnp.abs(acceleration) < tol, + 0, + acceleration, + ) + return acceleration + + def update_nodal_velocity( + self, elements, particles: _ParticlesState, dt: float, *args + ): + """Update the nodal momentum based on total force on nodes.""" + total_force = ( + elements.nodes.f_int + elements.nodes.f_ext + elements.nodes.f_damp + ) + acceleration = jnp.nan_to_num(jnp.divide(total_force, elements.nodes.mass)) + + velocity = elements.nodes.velocity + acceleration * dt + if elements.constraints: + velocity = self._apply_boundary_constraints_vel(elements, velocity) + velocity = jnp.where( + jnp.abs(velocity) < 1e-12, + 0, + velocity, + ) + return velocity, "velocity" + + @classmethod + def _update_nodal_velocity(cls, total_force, nvel, nmass, constraints, dt, tol): + """Update the nodal momentum based on total force on nodes.""" + acceleration = jnp.nan_to_num(jnp.divide(total_force, nmass)) + + velocity = nvel + acceleration * dt + if constraints: + velocity = cls._apply_boundary_constraints_vel(constraints, velocity) + velocity = jnp.where( + jnp.abs(velocity) < tol, + 0, + velocity, + ) + return velocity + + def update_nodal_momentum( + self, elements, particles: _ParticlesState, dt: float, *args + ): + """Update the nodal momentum based on total force on nodes.""" + momentum = elements.nodes.momentum.at[:].set( + elements.nodes.mass * elements.nodes.velocity + ) + momentum = jnp.where( + jnp.abs(momentum) < 1e-12, + 0, + momentum, + ) + return momentum, "momentum" + + @classmethod + def _update_nodal_momentum(cls, nmass, nvel, constraints, tol): + """Update the nodal momentum based on total force on nodes.""" + momentum = nmass * nvel + momentum = jnp.where( + jnp.abs(momentum) < tol, + 0, + momentum, + ) + return momentum + + @classmethod + def _apply_boundary_constraints_vel(cls, constraints, vel, *args): + """Apply boundary conditions for nodal velocity.""" + + # This assumes that the constraints don't have overlapping + # conditions. In case it does, only the first constraint will + # be applied. + def _func2(constraint, *, orig): + return constraint[1].apply_vel(orig, constraint[0]) + + partial_func = partial(_func2, orig=vel) + _out = tree_map( + partial_func, constraints, is_leaf=lambda x: isinstance(x, tuple) + ) + + def _f(x, *, orig): + return jnp.where(x == orig, jnp.nan, x) + + _pf = partial(_f, orig=vel) + _step_1 = tree_map(_pf, _out) + vel = tree_reduce( + lambda x, y: jnp.where(jnp.isnan(y), x, y), + [vel, _step_1], + ) + return vel + + @classmethod + def _apply_boundary_constraints_mom(cls, constraints, mom, mass, *args): + """Apply boundary conditions for nodal momentum.""" + + # This assumes that the constraints don't have overlapping + # conditions. In case it does, only the first constraint will + # be applied. + def _func2(constraint, *, mom, mass): + return constraint[1].apply_mom(mom, mass, constraint[0]) + + partial_func = partial(_func2, mom=mom, mass=mass) + _out = tree_map( + partial_func, constraints, is_leaf=lambda x: isinstance(x, tuple) + ) + + def _f(x, *, orig): + return jnp.where(x == orig, jnp.nan, x) + + _pf = partial(_f, orig=mom) + _step_1 = tree_map(_pf, _out) + mom = tree_reduce( + lambda x, y: jnp.where(jnp.isnan(y), x, y), + [mom, _step_1], + ) + return mom + + @classmethod + def _apply_boundary_constraints_acc(cls, constraints, orig, *args): + """Apply boundary conditions for nodal acceleration.""" + + # This assumes that the constraints don't have overlapping + # conditions. In case it does, only the first constraint will + # be applied. + def _func2(constraint, *, orig): + return constraint[1].apply_acc(orig, constraint[0]) + + partial_func = partial(_func2, orig=orig) + _out = tree_map( + partial_func, constraints, is_leaf=lambda x: isinstance(x, tuple) + ) + + def _f(x, *, orig): + return jnp.where(x == orig, jnp.nan, x) + + _pf = partial(_f, orig=orig) + _step_1 = tree_map(_pf, _out) + acc = tree_reduce( + lambda x, y: jnp.where(jnp.isnan(y), x, y), + [orig, _step_1], + ) + return acc + + @classmethod + def _apply_boundary_constraints(cls, nvel, nmom, nacc, nmass, constraints, *args): + if constraints: + vel = cls._apply_boundary_constraints_vel(constraints, nvel, *args) + mom = cls._apply_boundary_constraints_mom(constraints, nmom, nmass, *args) + acc = cls._apply_boundary_constraints_acc(constraints, nacc, *args) + return vel, mom, acc + + def apply_boundary_constraints(self, elements, *args): + if elements.constraints: + vel = self._apply_boundary_constraints_vel( + elements, elements.nodes.velocity, *args + ) + mom = self._apply_boundary_constraints_mom( + elements, elements.nodes.momentum, elements.nodes.mass, *args + ) + acc = self._apply_boundary_constraints_acc( + elements, elements.nodes.acceleration, *args + ) + + return elements.nodes.replace(velocity=vel, momentum=mom, acceleration=acc) - def compute_volume(self, *args): + def compute_volume(self, elements, *args): """Compute volume of all elements.""" - a = c = self.el_len[1] - b = d = self.el_len[0] + a = c = elements.el_len[1] + b = d = elements.el_len[0] p = q = jnp.sqrt(a**2 + b**2) vol = 0.25 * jnp.sqrt(4 * p * p * q * q - (a * a + c * c - b * b - d * d) ** 2) - self.volume = self.volume.at[:].set(vol) + volume = jnp.ones_like(elements.volume) * vol + return elements.replace(volume=volume) diff --git a/diffmpm/explicit.py b/diffmpm/explicit.py new file mode 100644 index 0000000..de8bea2 --- /dev/null +++ b/diffmpm/explicit.py @@ -0,0 +1,447 @@ +import abc +from dataclasses import dataclass +from functools import partial +from typing import Callable, NamedTuple, Optional + +from jax import Array, vmap +from jax.tree_util import tree_map, tree_structure, tree_transpose, tree_reduce + +from diffmpm.element import Quad4N, _ElementsState +from diffmpm.node import _reset_node_props +from diffmpm.particle import ( + _get_natural_coords, + _ParticlesState, + _compute_strain, + _update_particle_volume, + _compute_stress, + _update_particle_position_velocity, +) + + +class MeshState(NamedTuple): + elements: _ElementsState + particles: _ParticlesState + + +class Solver(abc.ABC): + @abc.abstractmethod + def init_state(*args, **kwargs): + pass + + @abc.abstractmethod + def update(*args, **kwargs): + pass + + def run(*args, **kwargs): + pass + + +def _reduce_attr(state_1, state_2, *, orig): + new_val = state_1 + state_2 - orig + return new_val + + +def _tree_transpose(pytree): + _out = tree_transpose( + tree_structure([0 for e in pytree]), tree_structure(pytree[0]), pytree + ) + return _out + + +@dataclass(eq=False) +class ExplicitSolver(Solver): + el_type: Quad4N + tol: float + dt: float + sim_steps: int + out_steps: int + out_dir: str + gravity: Array + scheme: str = "usf" + velocity_update: Optional[bool] = False + writer_func: Optional[Callable] = None + + def init_state(self, config): + elements = config["elements"] + particles = config["particles"] + return MeshState(elements=elements, particles=particles) + + def update(self, state: MeshState, step, *args, **kwargs): + _elements, _particles = state.elements, state.particles + # Nodal properties that are to be reset at the beginning of the + # update step. + new_nmass, new_nmom, new_nfint, new_nfext, new_nfdamp = _reset_node_props( + _elements.nodes + ) + + # New Element IDs for particles in each particle set. + # This is a `tree_map` function so that each particle set gets + # new EIDs. + new_peids: list = tree_map( + self.el_type._get_particles_element_ids, + _particles, + [_elements] * len(_particles), + is_leaf=lambda x: isinstance(x, _ParticlesState) + or isinstance(x, _ElementsState), + ) + # TODO: What if I calculate the mapped nodes for new eids and store + # here in a variable? This will allow reusing these without recomputing + # in every function. + + # new_pmapped_node_ids = vmap(self.el_type._get_mapped_nodes, (0, None))( + # new_peids, _elements + # ) + map_fn = vmap(self.el_type._get_mapped_nodes, (0, None)) + new_pmapped_node_ids = tree_map( + map_fn, + new_peids, + [_elements.nelements[0]] * len(_particles), + is_leaf=lambda x: isinstance(x, _ParticlesState) + or isinstance(x, _ElementsState), + ) + + # New natural coordinates of the particles. + # This is again a `tree_map`-ed function for each particle set. + # The signature of the function is + # `get_natural_coords(particles.loc, elements)` + # Attributes required: + # - Element IDs of the particles + # - Nodal coords of the elements corresponding to the above element ids. + def _leaf_fn(x): + return isinstance(x, _ParticlesState) or isinstance(x, Array) + + new_pxi = tree_map( + _get_natural_coords, + _particles, + new_pmapped_node_ids, + [_elements.nodes.loc] * len(_particles), + is_leaf=_leaf_fn, + ) + + # New nodal mass based on particle mass + # Required: + # - Nodal mass (new_nmass) + # - Particle natural coords (new_pxi) + # - Mapped nodes + # - Particle element IDs (new_peids) (list) + # new_nmass = self.el_type._compute_nodal_mass(new_nmass, new_pxi, new_peids) + temp_nmass = tree_map( + self.el_type._compute_nodal_mass, + [new_nmass] * len(_particles), + [p.mass for p in _particles], + new_pxi, + new_peids, + new_pmapped_node_ids, + [p.nparticles for p in _particles], + is_leaf=lambda x: isinstance(x, _ParticlesState) + or isinstance(x, Array) + or isinstance(x, int), + ) + partial_reduce_attr = partial(_reduce_attr, orig=new_nmass) + new_nmass = tree_reduce(partial_reduce_attr, temp_nmass) + + # New nodal momentum based on particle momentum + # Required: + # - Nodal momentum (new_nmom) + # - Particle natural coords (new_pxi) + # - Mapped nodes + # - Particle element IDs (new_peids) (list) + # new_nmom = _compute_nodal_momentum(new_nmom, new_xi, new_peids) + temp_nmom = tree_map( + self.el_type._compute_nodal_momentum, + [new_nmom] * len(_particles), + [p.mass for p in _particles], + [p.velocity for p in _particles], + new_pxi, + new_peids, + new_pmapped_node_ids, + [p.nparticles for p in _particles], + is_leaf=lambda x: isinstance(x, _ParticlesState) + or isinstance(x, Array) + or isinstance(x, int), + ) + partial_reduce_attr = partial(_reduce_attr, orig=new_nmom) + new_nmom = tree_reduce(partial_reduce_attr, temp_nmom) + + # New nodal velocity based on nodal momentum + # Required: + # - Nodal mass (new_nmass) + # - Current nodal velocity (_elements.nodes.velocity) + # - Nodal momentum (new_nmom) + # - Tolerance (tol) + # new_nvel = _compute_nodal_velocity( + # new_nmass, new_nmom, _elements.nodes.velocity, self.tol + # ) + temp_nvel = tree_map( + self.el_type._compute_nodal_velocity, + new_nmass, + new_nmom, + _elements.nodes.velocity, + ) + partial_reduce_attr = partial(_reduce_attr, orig=_elements.nodes.velocity) + new_nvel = tree_reduce(partial_reduce_attr, temp_nvel) + + # Apply boundary constraints on newly calculated props. + # Since nodal acceleration hasn't been updated yet, we + # use the current states nodal acceleration. + # Required: + # - Constraints (_elements.constraints) + # - Nodal velocity (new_nvel) + # - Nodal momentum (new_nmom) + # - Nodal acceleration (new_nacc) + new_nvel, new_nmom, new_nacc = self.el_type._apply_boundary_constraints( + new_nvel, + new_nmom, + _elements.nodes.acceleration, + new_nmass, + _elements.constraints, + ) + + if self.scheme == "usf": + # Compute particle strain + # Required: + # - Mapped node ids + # - Mapped node locs + # - Mapped node vels + # - Particle natural coords (new_pxi) + # - Current particle strains (_particles.strain) + # - Particles locs + # - Particle volumetric strains (_particles.volumetric_strain_centroid) + # ( + # new_pstrain_rate, + # new_pdstrain, + # new_pstrain, + # new_pdvolumetric_strain, + # new_pvolumetric_strain_centroid, + # ) = _compute_strain(_elements, _particles) + + _temp = tree_map( + _compute_strain, + [p.strain for p in _particles], + new_pxi, + [p.loc for p in _particles], + [p.volumetric_strain_centroid for p in _particles], + [p.nparticles for p in _particles], + new_pmapped_node_ids, + [_elements.nodes.loc] * len(_particles), + [_elements.nodes.velocity] * len(_particles), + [self.el_type] * len(_particles), + [self.dt] * len(_particles), + is_leaf=lambda x: isinstance(x, _ParticlesState) + or isinstance(x, Quad4N) + or isinstance(x, Array) + or isinstance(x, float), + ) + + _strains = _tree_transpose(_temp) + new_pstrain_rate = _strains["strain_rate"] + new_pdstrain = _strains["dstrain"] + new_pstrain = _strains["strain"] + new_pdvolumetric_strain = _strains["dvolumetric_strain"] + new_pvolumetric_strain_centroid = _strains["volumetric_strain_centroid"] + + # Compute new particle volumes based on updated strain + # Required: + # - Particle volumetric dstrain (new_pdvolumetric_strain) + # new_pvol, new_pdensity = _update_particle_volume(new_pdvolumetric_strain) + _temp = tree_map( + _update_particle_volume, + [p.volume for p in _particles], + [p.density for p in _particles], + new_pdvolumetric_strain, + ) + + new_pvol, new_pdensity = _tree_transpose(_temp) + # Compute particle stress + # Required: + # - Particle state since different materials need different + # particle properties to calculate stress. + # new_pstress = _compute_stress(_particles) + new_pstress = tree_map( + _compute_stress, + _particles, + is_leaf=lambda x: isinstance(x, _ParticlesState), + ) + + # Compute external forces on nodes + # Required: + # - Nodal external forces (new_nfext) + # - Particle natural coords (new_pxi) + # - Mapped Node ids + # new_nfext = self.el_type._compute_external_force(new_nfext, new_pxi, *args) + temp_nfext = tree_map( + self.el_type._compute_external_force, + [new_nfext] * len(_particles), + [p.f_ext for p in _particles], + new_pxi, + [p.nparticles for p in _particles], + new_pmapped_node_ids, + ) + partial_reduce_attr = partial(_reduce_attr, orig=new_nfext) + new_nfext = tree_reduce(partial_reduce_attr, temp_nfext) + + # Compute body forces on nodes + # Required: + # - Nodal external forces (new_nfext) + # - Particle natural coords (new_pxi) + # - Mapped Node ids + # - gravity + # new_nfext = _compute_body_force(new_nfext, new_pxi, gravity, *args) + temp_nfext = tree_map( + self.el_type._compute_body_force, + [new_nfext] * len(_particles), + [p.mass for p in _particles], + new_pxi, + new_pmapped_node_ids, + [p.nparticles for p in _particles], + [self.gravity] * len(_particles), + ) + partial_reduce_attr = partial(_reduce_attr, orig=new_nfext) + new_nfext = tree_reduce(partial_reduce_attr, temp_nfext) + + # TODO: Apply traction on particles + + # Apply nodal concentrated forces + # Required: + # - Concentrated forces on nodes (_elements.concentrated_nodal_forces) + # - Nodal external forces (new_nfext) + # - current time + # new_nfext = _apply_concentrated_nodal_forces( + # new_nfext, _elements.cnf, curr_time + # ) + # if _elements.concentrated_nodal_forces: + # temp_nfext = tree_map( + # self.el_type._apply_concentrated_nodal_forces, + # [new_nfext] * len(_elements.concentrated_nodal_forces), + # _elements.concentrated_nodal_forces, + # [self.dt] * len(_elements.concentrated_nodal_forces), + # is_leaf=lambda x: isinstance(x, NodalForce) + # or isinstance(x, float) + # or isinstance(x, Array), + # ) + # partial_reduce_attr = partial(_reduce_attr, orig=new_nfext) + # new_nfext = tree_reduce(partial_reduce_attr, temp_nfext) + + if _elements.concentrated_nodal_forces: + new_nfext = self.el_type._apply_concentrated_nodal_forces( + new_nfext, _elements.concentrated_nodal_forces, self.dt * step + ) + # Compute internal forces on nodes + # Required: + # - Mapped node ids + # - Mapped node locs + # - Nodal internal forces (new_nfint) + # - Particle natural coords (new_pxi) + # - Particle volume (new_pvol) + # - Particle stress (new_pstress) + temp_nfint = tree_map( + self.el_type._compute_internal_force, + [new_nfint] * len(_particles), + [_elements.nodes.loc] * len(_particles), + new_pmapped_node_ids, + new_pxi, + new_pvol, + [p.stress for p in _particles], + [p.nparticles for p in _particles], + ) + partial_reduce_attr = partial(_reduce_attr, orig=new_nfint) + new_nfint = tree_reduce(partial_reduce_attr, temp_nfint) + + if self.scheme == "usl": + # TODO: Calculate strains and stresses + pass + + # Update nodal acceleration based on nodal forces + # Required: + # - Nodal forces (new_nfint, new_nfext, new_nfdamp) + # - Nodal mass + # - Constraints (_elements.constraints) + # - Tolerance (self.tol) + total_force = new_nfint + new_nfext + new_nfdamp + new_nacc = self.el_type._update_nodal_acceleration( + total_force, new_nacc, new_nmass, _elements.constraints, self.tol + ) + # Update nodal acceleration based on nodal forces + # Required: + # - Nodal forces (new_nfint, new_nfext, new_nfdamp) + # - Nodal mass + # - Constraints (_elements.constraints) + # - Tolerance (self.tol) + new_nvel = self.el_type._update_nodal_velocity( + total_force, new_nvel, new_nmass, _elements.constraints, self.dt, self.tol + ) + + # Update nodal momentum based on nodal forces + # Required: + # - Nodal mass (new_nmass) + # - Nodal velocity (new_nvel) + # - Tolerance (self.tol) + new_nmom = self.el_type._update_nodal_momentum( + new_nmass, new_nvel, _elements.constraints, self.tol + ) + + # Update particle position and velocity + # Required: + # - Particle natural coords (new_pxi) + # - Timestep (self.dt) + # - self.velocity_update + # - Mapped node ids + # - Mapped node vels + # - Mapped node accelerations + # - Particle locs + _temp_new_vals = tree_map( + _update_particle_position_velocity, + [self.el_type] * len(_particles), + [p.loc for p in _particles], + [p.velocity for p in _particles], + [p.momentum for p in _particles], + [p.mass for p in _particles], + new_pxi, + new_pmapped_node_ids, + [new_nvel] * len(_particles), + [new_nacc] * len(_particles), + [self.velocity_update] * len(_particles), + [self.dt] * len(_particles), + is_leaf=lambda x: isinstance(x, Array) + or isinstance(x, Quad4N) + or isinstance(x, float) + or isinstance(x, bool), + ) + _new_vals = _tree_transpose(_temp_new_vals) + new_pvel = _new_vals["velocity"] + new_ploc = _new_vals["loc"] + new_pmom = _new_vals["momentum"] + + new_node_state = _elements.nodes.replace( + velocity=new_nvel, + acceleration=new_nacc, + mass=new_nmass, + momentum=new_nmom, + f_int=new_nfint, + f_ext=new_nfext, + f_damp=new_nfdamp, + ) + new_element_state = _elements.replace(nodes=new_node_state) + new_particle_states = [ + _p.replace( + loc=new_ploc, + element_ids=new_peids, + density=new_pdensity, + volume=new_pvol, + velocity=new_pvel, + momentum=new_pmom, + strain=new_pstrain, + stress=new_pstress, + strain_rate=new_pstrain_rate, + dstrain=new_pdstrain, + reference_loc=new_pxi, + dvolumetric_strain=new_pdvolumetric_strain, + volumetric_strain_centroid=new_pvolumetric_strain_centroid, + ) + for _p in _particles + ] + + new_mesh_state = MeshState( + elements=new_element_state, particles=new_particle_states + ) + return new_mesh_state diff --git a/diffmpm/io.py b/diffmpm/io.py index 822e50c..840dc69 100644 --- a/diffmpm/io.py +++ b/diffmpm/io.py @@ -154,9 +154,11 @@ def _parse_mesh(self, config): ] if config["mesh"]["type"] == "generator": - elements = element_cls( + total_elements = jnp.prod(jnp.array(config["mesh"]["nelements"])) + elementor = element_cls(total_elements=total_elements) + elements = elementor.init_state( config["mesh"]["nelements"], - jnp.prod(jnp.array(config["mesh"]["nelements"])), + total_elements, config["mesh"]["element_length"], constraints, concentrated_nodal_forces=self.parsed_config["external_loading"][ @@ -169,5 +171,6 @@ def _parse_mesh(self, config): ) self.parsed_config["elements"] = elements + self.parsed_config["elementor"] = elementor mesh = mesh_cls(self.parsed_config) return mesh diff --git a/diffmpm/mesh.py b/diffmpm/mesh.py index 2c06004..a20519a 100644 --- a/diffmpm/mesh.py +++ b/diffmpm/mesh.py @@ -6,7 +6,8 @@ from jax import lax, jit, tree_util from jax.tree_util import register_pytree_node_class, tree_map -from diffmpm.element import _Element +from diffmpm.element import _ElementState +import diffmpm.element as dfel from diffmpm.node import _NodesState from diffmpm.particle import _ParticlesState import diffmpm.particle as dpart @@ -29,8 +30,9 @@ class _MeshBase(abc.ABC): def __init__(self, config: dict): """Initialize mesh using configuration.""" self.particles: Sequence[_ParticlesState] = config["particles"] - self.elements: _Element = config["elements"] + self.elements: _ElementState = config["elements"] self.particle_tractions = config["particle_surface_traction"] + self.elementor = config["elementor"] # TODO: Change to allow called functions to return outputs def apply_on_elements(self, function: str, args: Tuple = ()): @@ -43,12 +45,12 @@ def apply_on_elements(self, function: str, args: Tuple = ()): args: tuple Parameters to be passed to the function. """ - f = getattr(self.elements, function) + f = getattr(self.elementor, function) - def _func(particles, *, func, fargs): - return func(particles, *fargs) + def _func(particles, *, func, elements, fargs): + return func(elements, particles, *fargs) - partial_func = partial(_func, func=f, fargs=args) + partial_func = partial(_func, func=f, elements=self.elements, fargs=args) _out = tree_map( partial_func, self.particles, @@ -58,6 +60,8 @@ def _func(particles, *, func, fargs): self.particles = _out elif function == "apply_boundary_constraints": self.elements.nodes = _out[0] + elif function == "compute_volume": + self.elements = _out[0] elif _out[0] is not None: _temp = tree_util.tree_transpose( tree_util.tree_structure([0 for e in _out]), @@ -96,7 +100,7 @@ def _func(particles, *, elements, fname, fargs): return f(particles, elements, *fargs) partial_func = partial( - _func, elements=self.elements, fname=function, fargs=args + _func, elements=self.elements, fname=function, fargs=(self.elementor, *args) ) new_states = tree_map( partial_func, @@ -155,7 +159,7 @@ def f(particles, *, ptraction, traction_val): def tree_flatten(self): children = (self.particles, self.elements) - aux_data = self.particle_tractions + aux_data = (self.elementor, self.particle_tractions) return (children, aux_data) @classmethod @@ -164,7 +168,8 @@ def tree_unflatten(cls, aux_data, children): { "particles": children[0], "elements": children[1], - "particle_surface_traction": aux_data, + "elementor": aux_data[0], + "particle_surface_traction": aux_data[1], } ) diff --git a/diffmpm/node.py b/diffmpm/node.py index eae41ba..815851b 100644 --- a/diffmpm/node.py +++ b/diffmpm/node.py @@ -61,16 +61,10 @@ def init_node_state( ) -def reset_node_state(state: _NodesState): - mass = state.mass.at[:].set(0) - momentum = state.momentum.at[:].set(0) - f_int = state.f_int.at[:].set(0) - f_ext = state.f_ext.at[:].set(0) - f_damp = state.f_damp.at[:].set(0) - return state.replace( - mass=mass, - momentum=momentum, - f_int=f_int, - f_ext=f_ext, - f_damp=f_damp, - ) +def _reset_node_props(state: _NodesState): + mass = jnp.zeros_like(state.mass) + momentum = jnp.zeros_like(state.momentum) + f_int = jnp.zeros_like(state.f_int) + f_ext = jnp.zeros_like(state.f_ext) + f_damp = jnp.zeros_like(state.f_damp) + return mass, momentum, f_int, f_ext, f_damp diff --git a/diffmpm/particle.py b/diffmpm/particle.py index 97921e0..8221468 100644 --- a/diffmpm/particle.py +++ b/diffmpm/particle.py @@ -2,10 +2,10 @@ import jax.numpy as jnp from jax import jit, lax, vmap -from jax.tree_util import register_pytree_node_class +from jax.tree_util import register_pytree_node_class, Partial from jax.typing import ArrayLike -from diffmpm.element import _Element +from diffmpm.element import _ElementsState from diffmpm.materials import _Material import chex @@ -138,21 +138,21 @@ def set_mass_volume(state, m: ArrayLike) -> _ParticlesState: return state.replace(mass=mass, volume=volume) -def compute_volume(state, elements: _Element, total_elements: int): +def compute_volume(state, elements: _ElementsState, elementor, total_elements: int): """Compute volume of all particles. Parameters ---------- state: Current state - elements: diffmpm._Element + 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`. """ particles_per_element = jnp.bincount( - state.element_ids, length=elements.total_elements + state.element_ids, length=elementor.total_elements ) vol = ( elements.volume.squeeze((1, 2))[state.element_ids] # type: ignore @@ -164,7 +164,7 @@ def compute_volume(state, elements: _Element, total_elements: int): return state.replace(mass=mass, size=size, volume=volume) -def update_natural_coords(state, elements: _Element): +def _get_natural_coords(particles, p_mapped_ids, eloc): r"""Update natural coordinates for the particles. Whenever the particles' physical coordinates change, their @@ -182,11 +182,40 @@ def update_natural_coords(state, elements: _Element): Parameters ---------- - elements: diffmpm.element._Element + elements: diffmpm.element._ElementState Elements based on which to update the natural coordinates of the particles. """ - t = vmap(jit(elements.id_to_node_loc))(state.element_ids) + t = eloc[p_mapped_ids.squeeze(-1)] + xi_coords = (particles.loc - (t[:, 0, ...] + t[:, 2, ...]) / 2) * ( + 2 / (t[:, 2, ...] - t[:, 0, ...]) + ) + return xi_coords + + +def update_natural_coords(state, elements: _ElementsState, elementor, *args): + r"""Update natural coordinates for the particles. + + Whenever the particles' physical coordinates change, their + natural coordinates need to be updated. This function updates + the natural coordinates of the particles based on the element + a particle is a part of. The update formula is + + \[ + \xi = (2x - (x_1^e + x_2^e)) / (x_2^e - x_1^e) + \] + + where \(x_i^e\) are the nodal coordinates of the element the + particle is in. If a particle is not in any element + (element_id = -1), its natural coordinate is set to 0. + + Parameters + ---------- + elements: diffmpm.element._ElementState + Elements based on which to update the natural coordinates + of the particles. + """ + t = vmap(Partial(elementor.id_to_node_loc, elements))(state.element_ids) xi_coords = (state.loc - (t[:, 0, ...] + t[:, 2, ...]) / 2) * ( 2 / (t[:, 2, ...] - t[:, 0, ...]) ) @@ -194,7 +223,7 @@ def update_natural_coords(state, elements: _Element): def update_position_velocity( - state, elements: _Element, dt: float, velocity_update: bool + state, elements: _ElementsState, elementor, dt: float, velocity_update: bool, *args ): """Transfer nodal velocity to particles and update particle position. @@ -202,7 +231,7 @@ def update_position_velocity( Parameters ---------- - elements: diffmpm.element._Element + elements: diffmpm.element._ElementState Elements whose nodes are used to transfer the velocity. dt: float Timestep. @@ -211,8 +240,10 @@ def update_position_velocity( velocity is calculated is interpolated nodal acceleration multiplied by dt. Default is False. """ - mapped_positions = elements.shapefn(state.reference_loc) - mapped_ids = vmap(jit(elements.id_to_node_ids))(state.element_ids).squeeze(-1) + mapped_positions = elementor.shapefn(state.reference_loc) + mapped_ids = vmap(Partial(elementor.id_to_node_ids, elements.nelements[0]))( + state.element_ids + ).squeeze(-1) nodal_velocity = jnp.sum( mapped_positions * elements.nodes.velocity[mapped_ids], axis=1 ) @@ -236,7 +267,56 @@ def update_position_velocity( return state.replace(velocity=velocity, loc=loc, momentum=momentum) -def _compute_strain_rate(state, dn_dx: ArrayLike, elements: _Element): +def _update_particle_position_velocity( + el_type, + ploc, + pvel, + pmom, + pmass, + pxi, + mapped_node_ids, + nvel, + nacc, + velocity_update, + dt, +): + """Transfer nodal velocity to particles and update particle position. + + The velocity is calculated based on the total force at nodes. + + Parameters + ---------- + elements: diffmpm.element._ElementState + Elements whose nodes are used to transfer the velocity. + dt: float + Timestep. + velocity_update: bool + If True, velocity is directly used as nodal velocity, else + velocity is calculated is interpolated nodal acceleration + multiplied by dt. Default is False. + """ + mapped_positions = el_type._shapefn(pxi) + mapped_ids = mapped_node_ids.squeeze(-1) + nodal_velocity = jnp.sum(mapped_positions * nvel[mapped_ids], axis=1) + nodal_acceleration = jnp.sum( + mapped_positions * nacc[mapped_ids], + axis=1, + ) + velocity = lax.cond( + velocity_update, + lambda sv, nv, na, t: nv, + lambda sv, nv, na, t: sv + na * t, + pvel, + nodal_velocity, + nodal_acceleration, + dt, + ) + loc = ploc.at[:].add(nodal_velocity * dt) + momentum = pmass * pvel + return {"velocity": velocity, "loc": loc, "momentum": momentum} + + +def _compute_strain_rate(mapped_vel, nparticles, dn_dx: ArrayLike): """Compute the strain rate for particles. Parameters @@ -244,15 +324,11 @@ def _compute_strain_rate(state, dn_dx: ArrayLike, elements: _Element): dn_dx: ArrayLike The gradient of the shape function. Expected shape `(nparticles, 1, ndim)` - elements: diffmpm.element._Element + elements: diffmpm.element._ElementState Elements whose nodes are used to calculate the strain rate. """ dn_dx = jnp.asarray(dn_dx) strain_rate = jnp.zeros((dn_dx.shape[0], 6, 1)) # (nparticles, 6, 1) - mapped_vel = vmap(jit(elements.id_to_node_vel))( - state.element_ids - ) # (nparticles, 2, 1) - temp = mapped_vel.squeeze(2) @jit @@ -265,14 +341,69 @@ def _step(pid, args): return dndx, nvel, strain_rate args = (dn_dx, temp, strain_rate) - _, _, strain_rate = lax.fori_loop(0, state.loc.shape[0], _step, args) + _, _, strain_rate = lax.fori_loop(0, nparticles, _step, args) strain_rate = jnp.where( jnp.abs(strain_rate) < 1e-12, jnp.zeros_like(strain_rate), strain_rate ) return strain_rate -def compute_strain(state, elements: _Element, dt: float): +def _compute_strain( + pstrain, + pxi, + ploc, + pvolumetric_strain_centroid, + nparticles, + mapped_node_ids, + nloc, + nvel, + el_type, + dt, +): + """Compute the strain on all particles. + + This is done by first calculating the strain rate for the particles + and then calculating strain as `strain += strain rate * dt`. + + Parameters + ---------- + elements: diffmpm.element._ElementState + Elements whose nodes are used to calculate the strain. + dt : float + Timestep. + """ + mapped_nodes = mapped_node_ids.squeeze(-1) + mapped_coords = nloc[mapped_nodes] + mapped_vel = nvel[mapped_nodes] + dn_dx_ = vmap(el_type._shapefn_grad)(pxi[:, jnp.newaxis, ...], mapped_coords) + new_strain_rate = _compute_strain_rate(mapped_vel, nparticles, dn_dx_) + new_dstrain = new_strain_rate * dt + + new_strain = pstrain + new_dstrain + centroids = jnp.zeros_like(ploc) + dn_dx_centroid_ = vmap(jit(el_type._shapefn_grad))( + centroids[:, jnp.newaxis, ...], mapped_coords + ) + strain_rate_centroid = _compute_strain_rate( + mapped_vel, + nparticles, + dn_dx_centroid_, + ) + ndim = ploc.shape[-1] + new_dvolumetric_strain = dt * strain_rate_centroid[:, :ndim].sum(axis=1) + new_volumetric_strain_centroid = ( + pvolumetric_strain_centroid + new_dvolumetric_strain + ) + return { + "strain_rate": new_strain_rate, + "dstrain": new_dstrain, + "strain": new_strain, + "dvolumetric_strain": new_dvolumetric_strain, + "volumetric_strain_centroid": new_volumetric_strain_centroid, + } + + +def compute_strain(state, elements: _ElementsState, elementor, dt: float, *args): """Compute the strain on all particles. This is done by first calculating the strain rate for the particles @@ -280,25 +411,29 @@ def compute_strain(state, elements: _Element, dt: float): Parameters ---------- - elements: diffmpm.element._Element + elements: diffmpm.element._ElementState Elements whose nodes are used to calculate the strain. dt : float Timestep. """ # breakpoint() - mapped_coords = vmap(jit(elements.id_to_node_loc))(state.element_ids).squeeze(2) - dn_dx_ = vmap(jit(elements.shapefn_grad))( + mapped_coords = vmap(Partial(elementor.id_to_node_loc, elements))( + state.element_ids + ).squeeze(2) + dn_dx_ = vmap(jit(elementor.shapefn_grad))( state.reference_loc[:, jnp.newaxis, ...], mapped_coords ) - strain_rate = _compute_strain_rate(state, dn_dx_, elements) + strain_rate = _compute_strain_rate(state, dn_dx_, elements, elementor) dstrain = state.dstrain.at[:].set(strain_rate * dt) strain = state.strain.at[:].add(dstrain) centroids = jnp.zeros_like(state.loc) - dn_dx_centroid_ = vmap(jit(elements.shapefn_grad))( + dn_dx_centroid_ = vmap(jit(elementor.shapefn_grad))( centroids[:, jnp.newaxis, ...], mapped_coords ) - strain_rate_centroid = _compute_strain_rate(state, dn_dx_centroid_, elements) + strain_rate_centroid = _compute_strain_rate( + state, dn_dx_centroid_, elements, elementor + ) ndim = state.loc.shape[-1] dvolumetric_strain = dt * strain_rate_centroid[:, :ndim].sum(axis=1) volumetric_strain_centroid = state.volumetric_strain_centroid.at[:].add( @@ -324,6 +459,17 @@ def compute_stress(state, *args): return state.replace(stress=stress) +def _compute_stress(state): + """Compute the strain on all particles. + + This calculation is governed by the material of the + particles. The stress calculated by the material is then + added to the particles current stress values. + """ + stress = state.stress.at[:].add(state.material.compute_stress(state)) + return stress + + def update_volume(state, *args): """Update volume based on central strain rate.""" volume = state.volume.at[:, 0, :].multiply(1 + state.dvolumetric_strain) @@ -331,6 +477,13 @@ def update_volume(state, *args): return state.replace(volume=volume, density=density) +def _update_particle_volume(pvol, pdensity, pdvolumetric_strain): + """Update volume based on central strain rate.""" + new_volume = pvol.at[:, 0, :].multiply(1 + pdvolumetric_strain) + new_density = pdensity.at[:, 0, :].divide(1 + pdvolumetric_strain) + return new_volume, new_density + + def assign_traction(state, pids: ArrayLike, dir: int, traction_: float): """Assign traction to particles. diff --git a/pyproject.toml b/pyproject.toml index ad356c2..6016ce1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,8 +12,8 @@ authors = [ readme = "README.md" version = "0.0.1" dependencies = [ - "jax[cpu]", - "click" + # "jax[cuda11_pip]", + # "click" ] classifiers = [ "Programming Language :: Python :: 3", @@ -26,3 +26,6 @@ mpm = "diffmpm.cli.mpm:mpm" [tool.black] line-length = 88 + +[tool.ruff] +line-length = 88 diff --git a/tests/test_particle.py b/tests/test_particle.py index 9c2acae..2dda962 100644 --- a/tests/test_particle.py +++ b/tests/test_particle.py @@ -1,7 +1,7 @@ import jax.numpy as jnp import pytest -from diffmpm.element import Quadrilateral4Node +from diffmpm.element import Quad4N from diffmpm.materials import init_simple from diffmpm.particle import init_particle_state import diffmpm.particle as dpar @@ -10,7 +10,7 @@ class TestParticles: @pytest.fixture def elements(self): - return Quadrilateral4Node((1, 1), 1, (1.0, 1.0), []) + return Quad4N().init_state(((1, 1), 1, (1.0, 1.0), [])) @pytest.fixture def particles(self): From 09bd261dc8855acbda180828f638bf8346c78eee Mon Sep 17 00:00:00 2001 From: Chahak Mehta Date: Mon, 31 Jul 2023 23:09:25 -0700 Subject: [PATCH 16/21] Working uniaxial stress and nodal forces benchmark --- diffmpm/element.py | 10 ++++ diffmpm/explicit.py | 85 +++++++++++++++++++++-------- diffmpm/io.py | 22 +++++--- diffmpm/materials/linear_elastic.py | 4 +- diffmpm/materials/simple.py | 4 +- diffmpm/mesh.py | 2 +- diffmpm/particle.py | 45 ++++++++++++--- tests/test_particle.py | 61 +++++++++++++++++++-- 8 files changed, 182 insertions(+), 51 deletions(-) diff --git a/diffmpm/element.py b/diffmpm/element.py index a068c57..b1dbb91 100644 --- a/diffmpm/element.py +++ b/diffmpm/element.py @@ -1183,6 +1183,16 @@ def apply_boundary_constraints(self, elements, *args): return elements.nodes.replace(velocity=vel, momentum=mom, acceleration=acc) + @classmethod + def _compute_volume(cls, el_len, evol): + """Compute volume of all elements.""" + a = c = el_len[1] + b = d = el_len[0] + p = q = jnp.sqrt(a**2 + b**2) + vol = 0.25 * jnp.sqrt(4 * p * p * q * q - (a * a + c * c - b * b - d * d) ** 2) + volume = jnp.ones_like(evol) * vol + return volume + def compute_volume(self, elements, *args): """Compute volume of all elements.""" a = c = elements.el_len[1] diff --git a/diffmpm/explicit.py b/diffmpm/explicit.py index de8bea2..399073e 100644 --- a/diffmpm/explicit.py +++ b/diffmpm/explicit.py @@ -1,4 +1,5 @@ import abc +import jax.numpy as jnp from dataclasses import dataclass from functools import partial from typing import Callable, NamedTuple, Optional @@ -15,6 +16,7 @@ _update_particle_volume, _compute_stress, _update_particle_position_velocity, + _compute_particle_volume, ) @@ -64,6 +66,36 @@ class ExplicitSolver(Solver): def init_state(self, config): elements = config["elements"] particles = config["particles"] + new_peids: list = tree_map( + self.el_type._get_particles_element_ids, + particles, + [elements] * len(particles), + is_leaf=lambda x: isinstance(x, _ParticlesState) + or isinstance(x, _ElementsState), + ) + new_evol = self.el_type._compute_volume(elements.el_len, elements.volume) + + temp_pprops = tree_map( + _compute_particle_volume, + new_peids, + [self.el_type.total_elements] * len(particles), + [new_evol] * len(particles), + [p.volume for p in particles], + [p.size for p in particles], + [p.mass for p in particles], + [p.density for p in particles], + ) + new_pprops = _tree_transpose(temp_pprops) + elements = elements.replace(volume=new_evol) + particles = [ + p.replace( + element_ids=new_peids, + mass=new_pprops["mass"][i], + size=new_pprops["size"][i], + volume=new_pprops["volume"][i], + ) + for i, p in enumerate(particles) + ] return MeshState(elements=elements, particles=particles) def update(self, state: MeshState, step, *args, **kwargs): @@ -74,6 +106,7 @@ def update(self, state: MeshState, step, *args, **kwargs): _elements.nodes ) + # breakpoint() # New Element IDs for particles in each particle set. # This is a `tree_map` function so that each particle set gets # new EIDs. @@ -84,13 +117,6 @@ def update(self, state: MeshState, step, *args, **kwargs): is_leaf=lambda x: isinstance(x, _ParticlesState) or isinstance(x, _ElementsState), ) - # TODO: What if I calculate the mapped nodes for new eids and store - # here in a variable? This will allow reusing these without recomputing - # in every function. - - # new_pmapped_node_ids = vmap(self.el_type._get_mapped_nodes, (0, None))( - # new_peids, _elements - # ) map_fn = vmap(self.el_type._get_mapped_nodes, (0, None)) new_pmapped_node_ids = tree_map( map_fn, @@ -118,6 +144,7 @@ def _leaf_fn(x): is_leaf=_leaf_fn, ) + # breakpoint() # New nodal mass based on particle mass # Required: # - Nodal mass (new_nmass) @@ -197,6 +224,7 @@ def _leaf_fn(x): _elements.constraints, ) + # breakpoint() if self.scheme == "usf": # Compute particle strain # Required: @@ -224,7 +252,7 @@ def _leaf_fn(x): [p.nparticles for p in _particles], new_pmapped_node_ids, [_elements.nodes.loc] * len(_particles), - [_elements.nodes.velocity] * len(_particles), + [new_nvel] * len(_particles), [self.el_type] * len(_particles), [self.dt] * len(_particles), is_leaf=lambda x: isinstance(x, _ParticlesState) @@ -251,6 +279,7 @@ def _leaf_fn(x): new_pdvolumetric_strain, ) + # breakpoint() new_pvol, new_pdensity = _tree_transpose(_temp) # Compute particle stress # Required: @@ -259,9 +288,14 @@ def _leaf_fn(x): # new_pstress = _compute_stress(_particles) new_pstress = tree_map( _compute_stress, - _particles, + [p.stress for p in _particles], + new_pstrain, + new_pdstrain, + [p.material for p in _particles], is_leaf=lambda x: isinstance(x, _ParticlesState), ) + # print(jnp.around(new_pstress[0].squeeze(), 5)) + # breakpoint() # Compute external forces on nodes # Required: @@ -341,12 +375,13 @@ def _leaf_fn(x): new_pmapped_node_ids, new_pxi, new_pvol, - [p.stress for p in _particles], + new_pstress, [p.nparticles for p in _particles], ) partial_reduce_attr = partial(_reduce_attr, orig=new_nfint) new_nfint = tree_reduce(partial_reduce_attr, temp_nfint) + # breakpoint() if self.scheme == "usl": # TODO: Calculate strains and stresses pass @@ -380,6 +415,7 @@ def _leaf_fn(x): new_nmass, new_nvel, _elements.constraints, self.tol ) + # breakpoint() # Update particle position and velocity # Required: # - Particle natural coords (new_pxi) @@ -411,6 +447,7 @@ def _leaf_fn(x): new_pvel = _new_vals["velocity"] new_ploc = _new_vals["loc"] new_pmom = _new_vals["momentum"] + # breakpoint() new_node_state = _elements.nodes.replace( velocity=new_nvel, @@ -424,21 +461,21 @@ def _leaf_fn(x): new_element_state = _elements.replace(nodes=new_node_state) new_particle_states = [ _p.replace( - loc=new_ploc, - element_ids=new_peids, - density=new_pdensity, - volume=new_pvol, - velocity=new_pvel, - momentum=new_pmom, - strain=new_pstrain, - stress=new_pstress, - strain_rate=new_pstrain_rate, - dstrain=new_pdstrain, - reference_loc=new_pxi, - dvolumetric_strain=new_pdvolumetric_strain, - volumetric_strain_centroid=new_pvolumetric_strain_centroid, + loc=new_ploc[i], + element_ids=new_peids[i], + density=new_pdensity[i], + volume=new_pvol[i], + velocity=new_pvel[i], + momentum=new_pmom[i], + strain=new_pstrain[i], + stress=new_pstress[i], + strain_rate=new_pstrain_rate[i], + dstrain=new_pdstrain[i], + reference_loc=new_pxi[i], + dvolumetric_strain=new_pdvolumetric_strain[i], + volumetric_strain_centroid=new_pvolumetric_strain_centroid[i], ) - for _p in _particles + for i, _p in enumerate(_particles) ] new_mesh_state = MeshState( diff --git a/diffmpm/io.py b/diffmpm/io.py index 840dc69..92330f8 100644 --- a/diffmpm/io.py +++ b/diffmpm/io.py @@ -5,16 +5,19 @@ from diffmpm import element as mpel from diffmpm import materials as mpmat -from diffmpm import mesh as mpmesh + +# from diffmpm import mesh as mpmesh from diffmpm.constraint import Constraint from diffmpm.forces import NodalForce, ParticleTraction from diffmpm.functions import Linear, Unit from diffmpm.particle import _ParticlesState, init_particle_state +from pathlib import Path class Config: def __init__(self, filepath): - self._filepath = filepath + self._filepath = Path(filepath).absolute() + self._basedir = self._filepath.parent self.parsed_config = {} self.parse() @@ -22,7 +25,9 @@ def parse(self): with open(self._filepath, "rb") as f: self._fileconfig = tl.load(f) - self.entity_sets = json.load(open(self._fileconfig["mesh"]["entity_sets"])) + self.entity_sets = json.load( + open(self._basedir.joinpath(self._fileconfig["mesh"]["entity_sets"])) + ) self._parse_meta(self._fileconfig) self._parse_output(self._fileconfig) self._parse_materials(self._fileconfig) @@ -31,7 +36,8 @@ def parse(self): self._parse_math_functions(self._fileconfig) self._parse_external_loading(self._fileconfig) mesh = self._parse_mesh(self._fileconfig) - return mesh + # return mesh + return self.parsed_config def _get_node_set_ids(self, set_ids): all_ids = [] @@ -68,7 +74,7 @@ def _parse_particles(self, config): particle_sets = [] for pset_config in config["particles"]: pmat = self.parsed_config["materials"][pset_config["material_id"]] - with open(pset_config["file"], "r") as f: + with open(self._basedir.joinpath(pset_config["file"]), "r") as f: ploc = jnp.asarray(json.load(f)) peids = jnp.zeros(ploc.shape[0], dtype=jnp.int32) pset = init_particle_state( @@ -141,7 +147,7 @@ def _parse_external_loading(self, config): def _parse_mesh(self, config): element_cls = getattr(mpel, config["mesh"]["element"]) - mesh_cls = getattr(mpmesh, f"Mesh{config['meta']['dimension']}D") + # mesh_cls = getattr(mpmesh, f"Mesh{config['meta']['dimension']}D") constraints = [] if "constraints" in config["mesh"]: @@ -172,5 +178,5 @@ def _parse_mesh(self, config): self.parsed_config["elements"] = elements self.parsed_config["elementor"] = elementor - mesh = mesh_cls(self.parsed_config) - return mesh + # mesh = mesh_cls(self.parsed_config) + # return mesh diff --git a/diffmpm/materials/linear_elastic.py b/diffmpm/materials/linear_elastic.py index 7e9c543..6cca66f 100644 --- a/diffmpm/materials/linear_elastic.py +++ b/diffmpm/materials/linear_elastic.py @@ -18,9 +18,9 @@ class _LinearElasticState: swave_velocity: float de: chex.ArrayDevice - def compute_stress(self, state): + def compute_stress(self, strain, dstrain): """Compute material stress.""" - dstress = self.de @ state.dstrain + dstress = self.de @ dstrain return dstress diff --git a/diffmpm/materials/simple.py b/diffmpm/materials/simple.py index 349190d..29f92ff 100644 --- a/diffmpm/materials/simple.py +++ b/diffmpm/materials/simple.py @@ -32,5 +32,5 @@ def __init__(self, material_properties): def __repr__(self): return f"SimpleMaterial(props={self.properties})" - def compute_stress(self, particles): - return particles.dstrain * self.properties["E"] + def compute_stress(self, strain, dstrain): + return dstrain * self.properties["E"] diff --git a/diffmpm/mesh.py b/diffmpm/mesh.py index a20519a..dde457b 100644 --- a/diffmpm/mesh.py +++ b/diffmpm/mesh.py @@ -6,7 +6,7 @@ from jax import lax, jit, tree_util from jax.tree_util import register_pytree_node_class, tree_map -from diffmpm.element import _ElementState +from diffmpm.element import _ElementsState import diffmpm.element as dfel from diffmpm.node import _NodesState from diffmpm.particle import _ParticlesState diff --git a/diffmpm/particle.py b/diffmpm/particle.py index 8221468..5302b48 100644 --- a/diffmpm/particle.py +++ b/diffmpm/particle.py @@ -138,6 +138,32 @@ def set_mass_volume(state, m: ArrayLike) -> _ParticlesState: return state.replace(mass=mass, volume=volume) +def _compute_particle_volume( + element_ids, total_elements, evol, pvol, psize, pmass, pdensity +): + """Compute volume of all particles. + + 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`. + """ + particles_per_element = jnp.bincount(element_ids, length=total_elements) + vol = ( + evol.squeeze((1, 2))[element_ids] # type: ignore + / particles_per_element[element_ids] + ) + volume = pvol.at[:, 0, 0].set(vol) + size = psize.at[:].set(volume ** (1 / psize.shape[-1])) + mass = pmass.at[:, 0, 0].set(vol * pdensity.squeeze()) + return {"mass": mass, "size": size, "volume": volume} + + def compute_volume(state, elements: _ElementsState, elementor, total_elements: int): """Compute volume of all particles. @@ -342,9 +368,7 @@ def _step(pid, args): args = (dn_dx, temp, strain_rate) _, _, strain_rate = lax.fori_loop(0, nparticles, _step, args) - strain_rate = jnp.where( - jnp.abs(strain_rate) < 1e-12, jnp.zeros_like(strain_rate), strain_rate - ) + strain_rate = jnp.where(jnp.abs(strain_rate) < 1e-12, 0, strain_rate) return strain_rate @@ -420,10 +444,12 @@ def compute_strain(state, elements: _ElementsState, elementor, dt: float, *args) mapped_coords = vmap(Partial(elementor.id_to_node_loc, elements))( state.element_ids ).squeeze(2) + mapped_vel = vmap(Partial(elementor.id_to_node_vel, elements))(state.element_ids) dn_dx_ = vmap(jit(elementor.shapefn_grad))( state.reference_loc[:, jnp.newaxis, ...], mapped_coords ) - strain_rate = _compute_strain_rate(state, dn_dx_, elements, elementor) + # strain_rate = _compute_strain_rate(state, dn_dx_, elements, elementor) + strain_rate = _compute_strain_rate(mapped_vel, state.nparticles, dn_dx_) dstrain = state.dstrain.at[:].set(strain_rate * dt) strain = state.strain.at[:].add(dstrain) @@ -431,8 +457,11 @@ def compute_strain(state, elements: _ElementsState, elementor, dt: float, *args) dn_dx_centroid_ = vmap(jit(elementor.shapefn_grad))( centroids[:, jnp.newaxis, ...], mapped_coords ) + # strain_rate_centroid = _compute_strain_rate( + # state, dn_dx_centroid_, elements, elementor + # ) strain_rate_centroid = _compute_strain_rate( - state, dn_dx_centroid_, elements, elementor + mapped_vel, state.nparticles, dn_dx_centroid_ ) ndim = state.loc.shape[-1] dvolumetric_strain = dt * strain_rate_centroid[:, :ndim].sum(axis=1) @@ -459,15 +488,15 @@ def compute_stress(state, *args): return state.replace(stress=stress) -def _compute_stress(state): +def _compute_stress(stress, strain, dstrain, material, *args): """Compute the strain on all particles. This calculation is governed by the material of the particles. The stress calculated by the material is then added to the particles current stress values. """ - stress = state.stress.at[:].add(state.material.compute_stress(state)) - return stress + new_stress = stress + material.compute_stress(strain, dstrain) + return new_stress def update_volume(state, *args): diff --git a/tests/test_particle.py b/tests/test_particle.py index 2dda962..1d4e2cd 100644 --- a/tests/test_particle.py +++ b/tests/test_particle.py @@ -1,5 +1,7 @@ import jax.numpy as jnp +from functools import partial import pytest +from jax import vmap from diffmpm.element import Quad4N from diffmpm.materials import init_simple @@ -8,9 +10,11 @@ class TestParticles: + elementor = Quad4N(total_elements=1) + @pytest.fixture def elements(self): - return Quad4N().init_state(((1, 1), 1, (1.0, 1.0), [])) + return self.elementor.init_state((1, 1), 1, (1.0, 1.0), []) @pytest.fixture def particles(self): @@ -26,22 +30,48 @@ def particles(self): ], ) def test_update_velocity(self, elements, particles, velocity_update, expected): - dpar.update_natural_coords(particles, elements) + dpar.update_natural_coords(particles, elements, self.elementor) elements.nodes = elements.nodes.replace( acceleration=elements.nodes.acceleration + 1 ) elements.nodes = elements.nodes.replace(velocity=elements.nodes.velocity + 1) + updated = dpar._update_particle_position_velocity( + Quad4N, + particles.loc, + particles.velocity, + particles.momentum, + particles.mass, + particles.reference_loc, + vmap(partial(self.elementor.id_to_node_ids, 1))(particles.element_ids), + elements.nodes.velocity, + elements.nodes.acceleration, + velocity_update, + 0.1, + ) particles = dpar.update_position_velocity( - particles, elements, 0.1, velocity_update + particles, elements, self.elementor, 0.1, velocity_update ) assert jnp.allclose(particles.velocity, expected) + assert jnp.allclose(updated["velocity"], expected) def test_compute_strain(self, elements, particles): elements.nodes = elements.nodes.replace( velocity=jnp.array([[0, 1], [0, 2], [0, 3], [0, 4]]).reshape(4, 1, 2) ) - particles = dpar.update_natural_coords(particles, elements) - particles = dpar.compute_strain(particles, elements, 0.1) + particles = dpar.update_natural_coords(particles, elements, self.elementor) + updated = dpar._compute_strain( + particles.strain, + particles.reference_loc, + particles.loc, + particles.volumetric_strain_centroid, + particles.nparticles, + vmap(partial(self.elementor.id_to_node_ids, 1))(particles.element_ids), + elements.nodes.loc, + elements.nodes.velocity, + Quad4N, + 0.1, + ) + particles = dpar.compute_strain(particles, elements, self.elementor, 0.1) assert jnp.allclose( particles.strain, jnp.array([[0, 0.2, 0, 0.1, 0, 0], [0, 0.2, 0, 0.1, 0, 0]]).reshape( @@ -49,10 +79,29 @@ def test_compute_strain(self, elements, particles): ), ) assert jnp.allclose(particles.volumetric_strain_centroid, jnp.array([0.2])) + assert jnp.allclose( + updated["strain"], + jnp.array([[0, 0.2, 0, 0.1, 0, 0], [0, 0.2, 0, 0.1, 0, 0]]).reshape( + 2, 6, 1 + ), + ) + assert jnp.allclose(updated["volumetric_strain_centroid"], jnp.array([0.2])) def test_compute_volume(self, elements, particles): - particles = dpar.compute_volume(particles, elements, elements.total_elements) + particles = dpar.compute_volume( + particles, elements, self.elementor, elements.total_elements + ) + props = dpar._compute_particle_volume( + particles.element_ids, + self.elementor.total_elements, + elements.volume, + particles.volume, + particles.size, + particles.mass, + particles.density, + ) assert jnp.allclose(particles.volume, jnp.array([0.5, 0.5]).reshape(2, 1, 1)) + assert jnp.allclose(props["volume"], jnp.array([0.5, 0.5]).reshape(2, 1, 1)) def test_assign_traction(self, elements, particles): particles = dpar.compute_volume(particles, elements, elements.total_elements) From a07cadf45150b77aa213b9d2a5c860082983af85 Mon Sep 17 00:00:00 2001 From: Chahak Mehta Date: Wed, 2 Aug 2023 19:22:15 -0700 Subject: [PATCH 17/21] Broken traction --- diffmpm/element.py | 29 ++++++++++ diffmpm/explicit.py | 127 +++++++++++++++++++++++++++++--------------- diffmpm/particle.py | 31 +++++++++++ 3 files changed, 145 insertions(+), 42 deletions(-) diff --git a/diffmpm/element.py b/diffmpm/element.py index b1dbb91..12bdbca 100644 --- a/diffmpm/element.py +++ b/diffmpm/element.py @@ -954,6 +954,35 @@ def _f(x, *, orig): f_ext = jnp.where(_step_2 == 0, nf_ext, _step_2) return f_ext + @classmethod + def _apply_particle_traction_forces( + cls, pxi, mapped_node_ids, nf_ext, ptraction, nparticles + ): + """Apply concentrated nodal forces. + + Parameters + ---------- + particles: Particles + Particles in the simulation. + """ + + @jit + def _step(pid, args): + f_ext, ptraction, mapped_pos, el_nodes = args + f_ext = f_ext.at[el_nodes[pid]].add(mapped_pos[pid] @ ptraction[pid]) + return f_ext, ptraction, mapped_pos, el_nodes + + mapped_positions = cls._shapefn(pxi) + mapped_nodes = mapped_node_ids.squeeze(-1) + args = ( + nf_ext, + ptraction, + mapped_positions, + mapped_nodes, + ) + f_ext, _, _, _ = lax.fori_loop(0, nparticles, _step, args) + return f_ext + def apply_particle_traction_forces(self, elements, particles: _ParticlesState): """Apply concentrated nodal forces. diff --git a/diffmpm/explicit.py b/diffmpm/explicit.py index 399073e..ecf9349 100644 --- a/diffmpm/explicit.py +++ b/diffmpm/explicit.py @@ -1,28 +1,44 @@ import abc -import jax.numpy as jnp from dataclasses import dataclass from functools import partial -from typing import Callable, NamedTuple, Optional +from typing import Callable, NamedTuple, Optional, Sequence from jax import Array, vmap -from jax.tree_util import tree_map, tree_structure, tree_transpose, tree_reduce +from jax.tree_util import tree_map, tree_reduce, tree_structure, tree_transpose from diffmpm.element import Quad4N, _ElementsState +from diffmpm.forces import ParticleTraction from diffmpm.node import _reset_node_props from diffmpm.particle import ( - _get_natural_coords, - _ParticlesState, + _assign_traction, + _compute_particle_volume, _compute_strain, - _update_particle_volume, _compute_stress, + _get_natural_coords, + _ParticlesState, _update_particle_position_velocity, - _compute_particle_volume, + _update_particle_volume, + _zero_traction, ) class MeshState(NamedTuple): elements: _ElementsState particles: _ParticlesState + particle_tractions: Sequence[ParticleTraction] + + @classmethod + def _apply_traction_on_particles( + cls, particles, particle_tractions, curr_time: float + ): + """Apply tractions on particles. + + Parameters + ---------- + curr_time: float + Current time in the simulation. + """ + pass class Solver(abc.ABC): @@ -96,7 +112,11 @@ def init_state(self, config): ) for i, p in enumerate(particles) ] - return MeshState(elements=elements, particles=particles) + return MeshState( + elements=elements, + particles=particles, + particle_tractions=config["particle_surface_traction"], + ) def update(self, state: MeshState, step, *args, **kwargs): _elements, _particles = state.elements, state.particles @@ -106,7 +126,6 @@ def update(self, state: MeshState, step, *args, **kwargs): _elements.nodes ) - # breakpoint() # New Element IDs for particles in each particle set. # This is a `tree_map` function so that each particle set gets # new EIDs. @@ -144,7 +163,6 @@ def _leaf_fn(x): is_leaf=_leaf_fn, ) - # breakpoint() # New nodal mass based on particle mass # Required: # - Nodal mass (new_nmass) @@ -224,7 +242,6 @@ def _leaf_fn(x): _elements.constraints, ) - # breakpoint() if self.scheme == "usf": # Compute particle strain # Required: @@ -235,14 +252,6 @@ def _leaf_fn(x): # - Current particle strains (_particles.strain) # - Particles locs # - Particle volumetric strains (_particles.volumetric_strain_centroid) - # ( - # new_pstrain_rate, - # new_pdstrain, - # new_pstrain, - # new_pdvolumetric_strain, - # new_pvolumetric_strain_centroid, - # ) = _compute_strain(_elements, _particles) - _temp = tree_map( _compute_strain, [p.strain for p in _particles], @@ -279,7 +288,6 @@ def _leaf_fn(x): new_pdvolumetric_strain, ) - # breakpoint() new_pvol, new_pdensity = _tree_transpose(_temp) # Compute particle stress # Required: @@ -294,8 +302,6 @@ def _leaf_fn(x): [p.material for p in _particles], is_leaf=lambda x: isinstance(x, _ParticlesState), ) - # print(jnp.around(new_pstress[0].squeeze(), 5)) - # breakpoint() # Compute external forces on nodes # Required: @@ -334,28 +340,66 @@ def _leaf_fn(x): new_nfext = tree_reduce(partial_reduce_attr, temp_nfext) # TODO: Apply traction on particles + new_ptraction = tree_map( + _zero_traction, + [p.traction for p in _particles], + is_leaf=lambda x: isinstance(x, _ParticlesState), + ) + + def func(ptract_, ptraction, pvol, psize, *, curr_time): + def f(ptraction, pvol, psize, *, ptract_, traction_val): + return _assign_traction( + ptraction, pvol, psize, ptract_.pids, ptract_.dir, traction_val + ) + + factor = ptract_.function.value(curr_time) + traction_val = factor * ptract_.traction + partial_f = partial(f, ptract_=ptract_, traction_val=traction_val) + traction_sets = tree_map( + partial_f, + ptraction, + pvol, + psize, + is_leaf=lambda x: isinstance(x, _ParticlesState), + ) + return tuple(traction_sets) + + partial_func = partial( + func, ptract_=state.particle_tractions, curr_time=step * self.dt + ) + if state.particle_tractions: + _out = tree_map( + partial_func, + state.particle_tractions, + new_ptraction, + new_pvol, + [p.size for p in _particles], + is_leaf=lambda x: isinstance(x, ParticleTraction) + or isinstance(x, Array), + ) + breakpoint() + _temp = _tree_transpose(_out) + new_ptraction = tree_reduce( + lambda x, y: x + y, _temp, is_leaf=lambda x: isinstance(x, list) + ) + + # breakpoint() + temp_nfext = tree_map( + self.el_type._apply_particle_traction_forces, + new_pxi, + new_pmapped_node_ids, + [new_nfext] * len(_particles), + new_ptraction, + [p.nparticles for p in _particles], + ) + partial_reduce_attr = partial(_reduce_attr, orig=new_nfext) + new_nfext = tree_reduce(partial_reduce_attr, temp_nfext) # Apply nodal concentrated forces # Required: # - Concentrated forces on nodes (_elements.concentrated_nodal_forces) # - Nodal external forces (new_nfext) # - current time - # new_nfext = _apply_concentrated_nodal_forces( - # new_nfext, _elements.cnf, curr_time - # ) - # if _elements.concentrated_nodal_forces: - # temp_nfext = tree_map( - # self.el_type._apply_concentrated_nodal_forces, - # [new_nfext] * len(_elements.concentrated_nodal_forces), - # _elements.concentrated_nodal_forces, - # [self.dt] * len(_elements.concentrated_nodal_forces), - # is_leaf=lambda x: isinstance(x, NodalForce) - # or isinstance(x, float) - # or isinstance(x, Array), - # ) - # partial_reduce_attr = partial(_reduce_attr, orig=new_nfext) - # new_nfext = tree_reduce(partial_reduce_attr, temp_nfext) - if _elements.concentrated_nodal_forces: new_nfext = self.el_type._apply_concentrated_nodal_forces( new_nfext, _elements.concentrated_nodal_forces, self.dt * step @@ -381,7 +425,6 @@ def _leaf_fn(x): partial_reduce_attr = partial(_reduce_attr, orig=new_nfint) new_nfint = tree_reduce(partial_reduce_attr, temp_nfint) - # breakpoint() if self.scheme == "usl": # TODO: Calculate strains and stresses pass @@ -415,7 +458,6 @@ def _leaf_fn(x): new_nmass, new_nvel, _elements.constraints, self.tol ) - # breakpoint() # Update particle position and velocity # Required: # - Particle natural coords (new_pxi) @@ -447,7 +489,6 @@ def _leaf_fn(x): new_pvel = _new_vals["velocity"] new_ploc = _new_vals["loc"] new_pmom = _new_vals["momentum"] - # breakpoint() new_node_state = _elements.nodes.replace( velocity=new_nvel, @@ -479,6 +520,8 @@ def _leaf_fn(x): ] new_mesh_state = MeshState( - elements=new_element_state, particles=new_particle_states + elements=new_element_state, + particles=new_particle_states, + particle_tractions=state.particle_tractions, ) return new_mesh_state diff --git a/diffmpm/particle.py b/diffmpm/particle.py index 5302b48..bdd0972 100644 --- a/diffmpm/particle.py +++ b/diffmpm/particle.py @@ -531,6 +531,37 @@ def assign_traction(state, pids: ArrayLike, dir: int, traction_: float): return traction +def _assign_traction( + ptraction, + pvol, + psize, + pids: ArrayLike, + dir: int, + traction_val_: float, +): + """Assign traction to particles. + + Parameters + ---------- + pids: ArrayLike + IDs of the particles to which traction should be applied. + dir: int + The direction in which traction should be applied. + traction_: float + Traction value to be applied in the direction. + """ + traction = ptraction.at[pids, 0, dir].add( + traction_val_ * pvol[pids, 0, 0] / psize[pids, 0, dir] + ) + return traction + + +def _zero_traction(traction): + """Set all traction values to 0.""" + traction = jnp.zeros_like(traction) + return traction + + def zero_traction(state, *args): """Set all traction values to 0.""" traction = state.traction.at[:].set(0) From b4b2e38b871f4e8728ab6f7aa95b3639563b5406 Mon Sep 17 00:00:00 2001 From: Chahak Mehta Date: Wed, 2 Aug 2023 20:28:56 -0700 Subject: [PATCH 18/21] Update tests --- tests/test_particle.py | 21 +++------------------ 1 file changed, 3 insertions(+), 18 deletions(-) diff --git a/tests/test_particle.py b/tests/test_particle.py index 1d4e2cd..6dd7b1d 100644 --- a/tests/test_particle.py +++ b/tests/test_particle.py @@ -48,10 +48,6 @@ def test_update_velocity(self, elements, particles, velocity_update, expected): velocity_update, 0.1, ) - particles = dpar.update_position_velocity( - particles, elements, self.elementor, 0.1, velocity_update - ) - assert jnp.allclose(particles.velocity, expected) assert jnp.allclose(updated["velocity"], expected) def test_compute_strain(self, elements, particles): @@ -71,14 +67,6 @@ def test_compute_strain(self, elements, particles): Quad4N, 0.1, ) - particles = dpar.compute_strain(particles, elements, self.elementor, 0.1) - assert jnp.allclose( - particles.strain, - jnp.array([[0, 0.2, 0, 0.1, 0, 0], [0, 0.2, 0, 0.1, 0, 0]]).reshape( - 2, 6, 1 - ), - ) - assert jnp.allclose(particles.volumetric_strain_centroid, jnp.array([0.2])) assert jnp.allclose( updated["strain"], jnp.array([[0, 0.2, 0, 0.1, 0, 0], [0, 0.2, 0, 0.1, 0, 0]]).reshape( @@ -88,9 +76,6 @@ def test_compute_strain(self, elements, particles): assert jnp.allclose(updated["volumetric_strain_centroid"], jnp.array([0.2])) def test_compute_volume(self, elements, particles): - particles = dpar.compute_volume( - particles, elements, self.elementor, elements.total_elements - ) props = dpar._compute_particle_volume( particles.element_ids, self.elementor.total_elements, @@ -100,9 +85,9 @@ def test_compute_volume(self, elements, particles): particles.mass, particles.density, ) - assert jnp.allclose(particles.volume, jnp.array([0.5, 0.5]).reshape(2, 1, 1)) assert jnp.allclose(props["volume"], jnp.array([0.5, 0.5]).reshape(2, 1, 1)) + @pytest.mark.skip() def test_assign_traction(self, elements, particles): particles = dpar.compute_volume(particles, elements, elements.total_elements) traction = dpar.assign_traction(particles, jnp.array([0]), 1, 10) @@ -112,5 +97,5 @@ def test_assign_traction(self, elements, particles): def test_zero_traction(self, particles): particles = particles.replace(traction=particles.traction + 1) - particles = dpar.zero_traction(particles) - assert jnp.all(particles.traction == 0) + traction = dpar._zero_traction(particles.traction) + assert jnp.all(traction == 0) From 1034d4f03c2d9e72277137befc3a58333047875a Mon Sep 17 00:00:00 2001 From: Chahak Mehta Date: Thu, 17 Aug 2023 12:31:17 -0400 Subject: [PATCH 19/21] Change loops to scan and add optim example --- diffmpm/element.py | 97 ++++++++++++++++++---- diffmpm/explicit.py | 14 ++-- diffmpm/functions.py | 2 +- diffmpm/materials/simple.py | 4 +- diffmpm/particle.py | 23 ++++-- optim_benchmark.py | 155 ++++++++++++++++++++++++++++++++++++ 6 files changed, 263 insertions(+), 32 deletions(-) create mode 100644 optim_benchmark.py diff --git a/diffmpm/element.py b/diffmpm/element.py index 12bdbca..6108a31 100644 --- a/diffmpm/element.py +++ b/diffmpm/element.py @@ -481,7 +481,7 @@ def _step(pid, args): @classmethod def _compute_internal_force( - cls, nf_int, nloc, mapped_node_ids, pxi, pvol, pstress, nparticles + cls, nf_int, nloc, mapped_node_ids, pxi, pvol, pstress, pids ): r"""Update the nodal internal force based on particle mass. @@ -528,6 +528,33 @@ def _step(pid, args): pstress, ) + def _scan_step(carry, pid): + ( + f_int, + pvol, + mapped_grads, + el_nodes, + pstress, + ) = carry + force = jnp.zeros((mapped_grads.shape[1], 1, 2)) + force = force.at[:, 0, 0].set( + mapped_grads[pid][:, 0] * pstress[pid][0] + + mapped_grads[pid][:, 1] * pstress[pid][3] + ) + force = force.at[:, 0, 1].set( + mapped_grads[pid][:, 1] * pstress[pid][1] + + mapped_grads[pid][:, 0] * pstress[pid][3] + ) + update = -pvol[pid] * force + f_int = f_int.at[el_nodes[pid]].add(update) + return ( + f_int, + pvol, + mapped_grads, + el_nodes, + pstress, + ), pid + # f_int = self.nodes.f_int.at[:].set(0) # f_int = elements.nodes.f_int mapped_nodes = mapped_node_ids.squeeze(-1) @@ -543,12 +570,14 @@ def _step(pid, args): mapped_nodes, pstress, ) - f_int, _, _, _, _ = lax.fori_loop(0, nparticles, _step, args) + # f_int, _, _, _, _ = lax.fori_loop(0, nparticles, _step, args) + final_carry, _ = lax.scan(_scan_step, args, pids) + f_int, _, _, _, _ = final_carry return f_int # Mapping from particles to nodes (P2G) @classmethod - def _compute_nodal_mass(cls, mass, pmass, pxi, peids, mapped_node_ids, nparticles): + def _compute_nodal_mass(cls, mass, pmass, pxi, peids, mapped_node_ids, pids): r"""Compute the nodal mass based on particle mass. The nodal mass is updated as a sum of particle mass for @@ -570,6 +599,11 @@ def _step(pid, args): mass = mass.at[el_nodes[pid]].add(pmass[pid] * mapped_pos[pid]) return pmass, mass, mapped_pos, el_nodes + def _scan_step(carry, pid): + pmass, mass, mapped_pos, el_nodes = carry + mass = mass.at[el_nodes[pid]].add(pmass[pid] * mapped_pos[pid]) + return (pmass, mass, mapped_pos, el_nodes), pid + mapped_positions = cls._shapefn(pxi) mapped_nodes = mapped_node_ids.squeeze(-1) args = ( @@ -578,7 +612,9 @@ def _step(pid, args): mapped_positions, mapped_nodes, ) - _, mass, _, _ = lax.fori_loop(0, nparticles, _step, args) + # _, mass, _, _ = lax.fori_loop(0, len(pids), _step, args) + final_carry, _ = lax.scan(_scan_step, args, pids) + _, mass, _, _ = final_carry return mass def compute_nodal_mass(self, elements, particles: _ParticlesState): @@ -621,7 +657,7 @@ def _step(pid, args): @classmethod def _compute_nodal_momentum( - cls, nmom, pmass, pvel, pxi, peids, mapped_node_ids, nparticles + cls, nmom, pmass, pvel, pxi, peids, mapped_node_ids, pids ): r"""Compute the nodal mass based on particle mass. @@ -644,8 +680,11 @@ def _step(pid, args): new_mom = mom.at[el_nodes[pid]].add(mapped_pos[pid] @ pmom[pid]) return pmom, new_mom, mapped_pos, el_nodes - # curr_mom = elements.nodes.momentum.at[:].set(0) - # curr_mom = elements.nodes.momentum + def _scan_step(carry, pid): + pmom, mom, mapped_pos, el_nodes = carry + new_mom = mom.at[el_nodes[pid]].add(mapped_pos[pid] @ pmom[pid]) + return (pmom, new_mom, mapped_pos, el_nodes), pid + mapped_nodes = mapped_node_ids.squeeze(-1) mapped_positions = cls._shapefn(pxi) args = ( @@ -654,9 +693,11 @@ def _step(pid, args): mapped_positions, mapped_nodes, ) - _, new_momentum, _, _ = lax.fori_loop(0, nparticles, _step, args) - new_momentum = jnp.where(jnp.abs(new_momentum) < 1e-12, 0, new_momentum) - return new_momentum + _, new_momentum, _, _ = lax.fori_loop(0, len(pids), _step, args) + final_carry, _ = lax.scan(_scan_step, args, pids) + _, new_nmom, _, _ = final_carry + new_nmom = jnp.where(jnp.abs(new_nmom) < 1e-12, 0, new_nmom) + return new_nmom def compute_nodal_momentum(self, elements, particles: _ParticlesState): r"""Compute the nodal mass based on particle mass. @@ -767,7 +808,7 @@ def _step(pid, args): return f_ext, "f_ext" @classmethod - def _compute_external_force(cls, f_ext, pf_ext, pxi, nparticles, mapped_node_ids): + def _compute_external_force(cls, f_ext, pf_ext, pxi, pids, mapped_node_ids): r"""Update the nodal external force based on particle f_ext. The nodal force is updated as a sum of particle external @@ -789,7 +830,11 @@ def _step(pid, args): f_ext = f_ext.at[el_nodes[pid]].add(mapped_pos[pid] @ pf_ext[pid]) return f_ext, pf_ext, mapped_pos, el_nodes - # f_ext = elements.nodes.f_ext.at[:].set(0) + def _scan_step(carry, pid): + f_ext, pf_ext, mapped_pos, el_nodes = carry + f_ext = f_ext.at[el_nodes[pid]].add(mapped_pos[pid] @ pf_ext[pid]) + return (f_ext, pf_ext, mapped_pos, el_nodes), pid + mapped_positions = cls._shapefn(pxi) mapped_nodes = mapped_node_ids.squeeze(-1) args = ( @@ -798,7 +843,9 @@ def _step(pid, args): mapped_positions, mapped_nodes, ) - f_ext, _, _, _ = lax.fori_loop(0, nparticles, _step, args) + # f_ext, _, _, _ = lax.fori_loop(0, len(pids), _step, args) + final_carry, _ = lax.scan(_scan_step, args, pids) + f_ext, _, _, _ = final_carry return f_ext def compute_body_force( @@ -844,7 +891,7 @@ def _step(pid, args): @classmethod def _compute_body_force( - cls, nf_ext, pmass, pxi, mapped_node_ids, nparticles, gravity: ArrayLike + cls, nf_ext, pmass, pxi, mapped_node_ids, pids, gravity: ArrayLike ): r"""Update the nodal external force based on particle mass. @@ -869,6 +916,13 @@ def _step(pid, args): ) return f_ext, pmass, mapped_pos, el_nodes, gravity + def _scan_step(carry, pid): + f_ext, pmass, mapped_pos, el_nodes, gravity = args + f_ext = f_ext.at[el_nodes[pid]].add( + mapped_pos[pid] @ (pmass[pid] * gravity) + ) + return (f_ext, pmass, mapped_pos, el_nodes, gravity), pid + mapped_positions = cls._shapefn(pxi) mapped_nodes = mapped_node_ids.squeeze(-1) args = ( @@ -878,7 +932,9 @@ def _step(pid, args): mapped_nodes, gravity, ) - f_ext, _, _, _, _ = lax.fori_loop(0, nparticles, _step, args) + # f_ext, _, _, _, _ = lax.fori_loop(0, nparticles, _step, args) + final_carry, _ = lax.scan(_scan_step, args, pids) + f_ext, _, _, _, _ = final_carry return f_ext def apply_concentrated_nodal_forces( @@ -956,7 +1012,7 @@ def _f(x, *, orig): @classmethod def _apply_particle_traction_forces( - cls, pxi, mapped_node_ids, nf_ext, ptraction, nparticles + cls, pxi, mapped_node_ids, nf_ext, ptraction, pids ): """Apply concentrated nodal forces. @@ -972,6 +1028,11 @@ def _step(pid, args): f_ext = f_ext.at[el_nodes[pid]].add(mapped_pos[pid] @ ptraction[pid]) return f_ext, ptraction, mapped_pos, el_nodes + def _scan_step(carry, pid): + f_ext, ptraction, mapped_pos, el_nodes = carry + f_ext = f_ext.at[el_nodes[pid]].add(mapped_pos[pid] @ ptraction[pid]) + return (f_ext, ptraction, mapped_pos, el_nodes), pid + mapped_positions = cls._shapefn(pxi) mapped_nodes = mapped_node_ids.squeeze(-1) args = ( @@ -980,7 +1041,9 @@ def _step(pid, args): mapped_positions, mapped_nodes, ) - f_ext, _, _, _ = lax.fori_loop(0, nparticles, _step, args) + # f_ext, _, _, _ = lax.fori_loop(0, nparticles, _step, args) + final_carry, _ = lax.scan(_scan_step, args, pids) + f_ext, _, _, _ = final_carry return f_ext def apply_particle_traction_forces(self, elements, particles: _ParticlesState): diff --git a/diffmpm/explicit.py b/diffmpm/explicit.py index ecf9349..800c8ef 100644 --- a/diffmpm/explicit.py +++ b/diffmpm/explicit.py @@ -177,7 +177,7 @@ def _leaf_fn(x): new_pxi, new_peids, new_pmapped_node_ids, - [p.nparticles for p in _particles], + [p.ids for p in _particles], is_leaf=lambda x: isinstance(x, _ParticlesState) or isinstance(x, Array) or isinstance(x, int), @@ -200,7 +200,7 @@ def _leaf_fn(x): new_pxi, new_peids, new_pmapped_node_ids, - [p.nparticles for p in _particles], + [p.ids for p in _particles], is_leaf=lambda x: isinstance(x, _ParticlesState) or isinstance(x, Array) or isinstance(x, int), @@ -258,7 +258,7 @@ def _leaf_fn(x): new_pxi, [p.loc for p in _particles], [p.volumetric_strain_centroid for p in _particles], - [p.nparticles for p in _particles], + [p.ids for p in _particles], new_pmapped_node_ids, [_elements.nodes.loc] * len(_particles), [new_nvel] * len(_particles), @@ -314,7 +314,7 @@ def _leaf_fn(x): [new_nfext] * len(_particles), [p.f_ext for p in _particles], new_pxi, - [p.nparticles for p in _particles], + [p.ids for p in _particles], new_pmapped_node_ids, ) partial_reduce_attr = partial(_reduce_attr, orig=new_nfext) @@ -333,7 +333,7 @@ def _leaf_fn(x): [p.mass for p in _particles], new_pxi, new_pmapped_node_ids, - [p.nparticles for p in _particles], + [p.ids for p in _particles], [self.gravity] * len(_particles), ) partial_reduce_attr = partial(_reduce_attr, orig=new_nfext) @@ -390,7 +390,7 @@ def f(ptraction, pvol, psize, *, ptract_, traction_val): new_pmapped_node_ids, [new_nfext] * len(_particles), new_ptraction, - [p.nparticles for p in _particles], + [p.ids for p in _particles], ) partial_reduce_attr = partial(_reduce_attr, orig=new_nfext) new_nfext = tree_reduce(partial_reduce_attr, temp_nfext) @@ -420,7 +420,7 @@ def f(ptraction, pvol, psize, *, ptract_, traction_val): new_pxi, new_pvol, new_pstress, - [p.nparticles for p in _particles], + [p.ids for p in _particles], ) partial_reduce_attr = partial(_reduce_attr, orig=new_nfint) new_nfint = tree_reduce(partial_reduce_attr, temp_nfint) diff --git a/diffmpm/functions.py b/diffmpm/functions.py index 90b55c4..ef6b3e1 100644 --- a/diffmpm/functions.py +++ b/diffmpm/functions.py @@ -22,7 +22,7 @@ def value(self, x): return 1.0 def tree_flatten(self): - return ((), (self.id)) + return ((), (self.id,)) @classmethod def tree_unflatten(cls, aux_data, children): diff --git a/diffmpm/materials/simple.py b/diffmpm/materials/simple.py index 29f92ff..142ead6 100644 --- a/diffmpm/materials/simple.py +++ b/diffmpm/materials/simple.py @@ -12,8 +12,8 @@ class _SimpleMaterialState: density: float state_vars: () - def compute_stress(self, particles): - return particles.dstrain * self.E + def compute_stress(self, strain, dstrain): + return dstrain * self.E def init_simple(material_properties): diff --git a/diffmpm/particle.py b/diffmpm/particle.py index bdd0972..b8cac78 100644 --- a/diffmpm/particle.py +++ b/diffmpm/particle.py @@ -13,6 +13,7 @@ @chex.dataclass(frozen=True) class _ParticlesState: + ids: chex.ArrayDevice nparticles: int loc: chex.ArrayDevice material: _Material @@ -68,6 +69,7 @@ def init_particle_state( f"`loc` should be of size (nparticles, 1, ndim); " f"found {loc.shape}" ) + ids = jnp.arange(loc.shape[0]) mass = jnp.ones((loc.shape[0], 1, 1)) density = jnp.ones_like(mass) * material.density volume = jnp.ones_like(mass) @@ -88,6 +90,7 @@ def init_particle_state( if material.state_vars: state_vars = material.initialize_state_variables(loc.shape[0]) return _ParticlesState( + ids=ids, nparticles=loc.shape[0], loc=loc, material=material, @@ -342,7 +345,7 @@ def _update_particle_position_velocity( return {"velocity": velocity, "loc": loc, "momentum": momentum} -def _compute_strain_rate(mapped_vel, nparticles, dn_dx: ArrayLike): +def _compute_strain_rate(mapped_vel, pids, dn_dx: ArrayLike): """Compute the strain rate for particles. Parameters @@ -366,8 +369,18 @@ def _step(pid, args): strain_rate = strain_rate.at[pid, 3].add(matmul[0, 1] + matmul[1, 0]) return dndx, nvel, strain_rate + def _scan_step(carry, pid): + dndx, nvel, strain_rate = carry + matmul = dndx[pid].T @ nvel[pid] + strain_rate = strain_rate.at[pid, 0].add(matmul[0, 0]) + strain_rate = strain_rate.at[pid, 1].add(matmul[1, 1]) + strain_rate = strain_rate.at[pid, 3].add(matmul[0, 1] + matmul[1, 0]) + return (dndx, nvel, strain_rate), pid + args = (dn_dx, temp, strain_rate) - _, _, strain_rate = lax.fori_loop(0, nparticles, _step, args) + # _, _, strain_rate = lax.fori_loop(0, nparticles, _step, args) + final_carry, _ = lax.scan(_scan_step, args, pids) + _, _, strain_rate = final_carry strain_rate = jnp.where(jnp.abs(strain_rate) < 1e-12, 0, strain_rate) return strain_rate @@ -377,7 +390,7 @@ def _compute_strain( pxi, ploc, pvolumetric_strain_centroid, - nparticles, + pids, mapped_node_ids, nloc, nvel, @@ -400,7 +413,7 @@ def _compute_strain( mapped_coords = nloc[mapped_nodes] mapped_vel = nvel[mapped_nodes] dn_dx_ = vmap(el_type._shapefn_grad)(pxi[:, jnp.newaxis, ...], mapped_coords) - new_strain_rate = _compute_strain_rate(mapped_vel, nparticles, dn_dx_) + new_strain_rate = _compute_strain_rate(mapped_vel, pids, dn_dx_) new_dstrain = new_strain_rate * dt new_strain = pstrain + new_dstrain @@ -410,7 +423,7 @@ def _compute_strain( ) strain_rate_centroid = _compute_strain_rate( mapped_vel, - nparticles, + pids, dn_dx_centroid_, ) ndim = ploc.shape[-1] diff --git a/optim_benchmark.py b/optim_benchmark.py new file mode 100644 index 0000000..6742626 --- /dev/null +++ b/optim_benchmark.py @@ -0,0 +1,155 @@ +from typing import NamedTuple +from functools import partial +import matplotlib.pyplot as plt + +import jax +import jax.numpy as jnp +import optax +from tqdm import tqdm + +from diffmpm.constraint import Constraint +from diffmpm.element import Quad4N, Quad4NState +from diffmpm.explicit import ExplicitSolver +from diffmpm.forces import NodalForce +from diffmpm.functions import Unit +from diffmpm.io import Config +from diffmpm.materials import init_simple, init_linear_elastic +from diffmpm.particle import _ParticlesState, init_particle_state + +jax.config.update("jax_platform_name", "cpu") + +config = Config("./benchmarks/2d/uniaxial_stress/mpm-uniaxial-stress.toml") +# config = Config("./benchmarks/2d/uniaxial_particle_traction/mpm-particle-traction.toml") +# config = Config("./benchmarks/2d/uniaxial_nodal_forces/mpm-nodal-forces.toml") +# config = Config("./benchmarks/2d/hydrostatic_column/mpm.toml") +# parsed_config = config.parse() +# cnf = [NodalForce(node_ids=jnp.array([0, 1]), function=Unit(-1), dir=1, force=1.5)] +# material = NamedTuple("Simple", density=1, E=1, state_vars={}) +# ploc = jnp.array([[0.5, 0.5], [0.5, 0.5]]).reshape(2, 1, 2) +# pmat = material(density=1.0, E=1.0, state_vars={}) +# pmat = init_simple({"density": 1, "E": 100, "id": 1}) +# peids = jnp.array([1]) +# particles = [init_particle_state(ploc, pmat, peids)] + +# cls = Quad4N(total_elements=1) +# elements = cls.init_state( +# (1, 1), +# 1, +# (1, 1), +# [(jnp.array([0]), Constraint(0, 2))], +# concentrated_nodal_forces=cnf, +# ) + +solver = ExplicitSolver( + el_type=config.parsed_config["elementor"], + tol=1e-12, + scheme=config.parsed_config["meta"]["scheme"], + dt=config.parsed_config["meta"]["dt"], + velocity_update=config.parsed_config["meta"]["velocity_update"], + sim_steps=config.parsed_config["meta"]["nsteps"], + out_steps=config.parsed_config["output"]["step_frequency"], + out_dir=config.parsed_config["output"]["format"], + gravity=config.parsed_config["external_loading"]["gravity"], +) + +init_vals = solver.init_state( + { + "elements": config.parsed_config["elements"], + "particles": config.parsed_config["particles"], + "particle_surface_traction": config.parsed_config["particle_surface_traction"], + } +) + +jit_updated = init_vals +jitted_update = jax.jit(solver.update) +for step in tqdm(range(20)): + # jit_updated = solver.update(jit_updated, step + 1) + jit_updated = jitted_update(jit_updated, step + 1) + +true_vel = jit_updated.particles[0].stress + + +def compute_loss(params, *, solver, target_vel, config): + # material = init_simple({"E": params, "density": 1, "id": -1}) + material = init_linear_elastic( + {"youngs_modulus": params, "density": 1, "poisson_ratio": 0, "id": -1} + ) + # breakpoint() + particles_ = [ + init_particle_state( + config.parsed_config["particles"][0].loc, + material, + config.parsed_config["particles"][0].element_ids, + init_vel=jnp.asarray([1.0, 0.0]), + ) + ] + init_vals = solver.init_state( + { + "elements": config.parsed_config["elements"], + "particles": particles_, + "particle_surface_traction": config.parsed_config[ + "particle_surface_traction" + ], + } + ) + result = init_vals + for step in tqdm(range(20), leave=False): + result = jitted_update(result, step + 1) + vel = result.particles[0].stress + loss = jnp.linalg.norm(vel - target_vel) + return loss + + +def optax_adam(params, niter, mpm, target_vel, config): + # Initialize parameters of the model + optimizer. + start_learning_rate = 1 + optimizer = optax.adam(start_learning_rate) + opt_state = optimizer.init(params) + + param_list = [] + loss_list = [] + # A simple update loop. + t = tqdm(range(niter), desc=f"E: {params}") + partial_f = partial(compute_loss, solver=mpm, target_vel=target_vel, config=config) + for _ in t: + lo, grads = jax.value_and_grad(partial_f, argnums=0)(params) + updates, opt_state = optimizer.update(grads, opt_state) + params = optax.apply_updates(params, updates) + t.set_description(f"YM: {params}") + param_list.append(params) + loss_list.append(lo) + return param_list, loss_list + + +params = 900.5 +# material = init_simple({"E": params, "density": 1, "id": -1}) +material = init_linear_elastic( + {"youngs_modulus": params, "density": 1, "poisson_ratio": 0, "id": -1} +) +particles = [ + init_particle_state( + config.parsed_config["particles"][0].loc, + material, + config.parsed_config["particles"][0].element_ids, + ) +] + +init_vals = solver.init_state( + { + "elements": config.parsed_config["elements"], + "particles": particles, + "particle_surface_traction": config.parsed_config["particle_surface_traction"], + } +) +param_list, loss_list = optax_adam( + params, 100, solver, true_vel, config +) # ADAM optimizer + +fig, ax = plt.subplots(1, 2, figsize=(16, 6)) +ax[0].plot(param_list, "ko", markersize=2, label="E") +ax[0].grid() +ax[0].legend() +ax[1].plot(loss_list, "ko", markersize=2, label="Loss") +ax[1].grid() +ax[1].legend() +plt.show() From 03f112ca6df94959e09f61a91acc5abb550edadf Mon Sep 17 00:00:00 2001 From: Chahak Mehta Date: Thu, 17 Aug 2023 13:03:24 -0400 Subject: [PATCH 20/21] Update optim file --- optim_benchmark.py | 31 +++++++++++++++++++++---------- 1 file changed, 21 insertions(+), 10 deletions(-) diff --git a/optim_benchmark.py b/optim_benchmark.py index 6742626..993ee4e 100644 --- a/optim_benchmark.py +++ b/optim_benchmark.py @@ -72,9 +72,13 @@ def compute_loss(params, *, solver, target_vel, config): # material = init_simple({"E": params, "density": 1, "id": -1}) material = init_linear_elastic( - {"youngs_modulus": params, "density": 1, "poisson_ratio": 0, "id": -1} + { + "youngs_modulus": params["ym"], + "density": 1, + "poisson_ratio": 0, + "id": -1, + } ) - # breakpoint() particles_ = [ init_particle_state( config.parsed_config["particles"][0].loc, @@ -102,11 +106,11 @@ def compute_loss(params, *, solver, target_vel, config): def optax_adam(params, niter, mpm, target_vel, config): # Initialize parameters of the model + optimizer. - start_learning_rate = 1 + start_learning_rate = 4 optimizer = optax.adam(start_learning_rate) opt_state = optimizer.init(params) - param_list = [] + param_list = {"ym": [], "pr": []} loss_list = [] # A simple update loop. t = tqdm(range(niter), desc=f"E: {params}") @@ -115,16 +119,23 @@ def optax_adam(params, niter, mpm, target_vel, config): lo, grads = jax.value_and_grad(partial_f, argnums=0)(params) updates, opt_state = optimizer.update(grads, opt_state) params = optax.apply_updates(params, updates) - t.set_description(f"YM: {params}") - param_list.append(params) + t.set_description(f"YM: {params['ym']:.2f}") + param_list["ym"].append(params["ym"]) + # param_list["pr"].append(params["pr"]) loss_list.append(lo) return param_list, loss_list -params = 900.5 +# params = {"pr": 0.4} +params = {"ym": 1101.0} # material = init_simple({"E": params, "density": 1, "id": -1}) material = init_linear_elastic( - {"youngs_modulus": params, "density": 1, "poisson_ratio": 0, "id": -1} + { + "youngs_modulus": params["ym"], + "density": 1, + "poisson_ratio": 0, + "id": -1, + } ) particles = [ init_particle_state( @@ -142,11 +153,11 @@ def optax_adam(params, niter, mpm, target_vel, config): } ) param_list, loss_list = optax_adam( - params, 100, solver, true_vel, config + params, 200, solver, true_vel, config ) # ADAM optimizer fig, ax = plt.subplots(1, 2, figsize=(16, 6)) -ax[0].plot(param_list, "ko", markersize=2, label="E") +ax[0].plot(param_list["ym"], "ko", markersize=2, label="E") ax[0].grid() ax[0].legend() ax[1].plot(loss_list, "ko", markersize=2, label="Loss") From fcec79623a5f73243094de995caddb2ed5a362f6 Mon Sep 17 00:00:00 2001 From: Chahak Mehta Date: Thu, 17 Aug 2023 13:14:29 -0400 Subject: [PATCH 21/21] Add uniaxial stress optim example plot --- .../optim_benchmark.py | 7 +++---- examples/optim_uniaxial_stress.png | Bin 0 -> 38565 bytes 2 files changed, 3 insertions(+), 4 deletions(-) rename optim_benchmark.py => examples/optim_benchmark.py (96%) create mode 100644 examples/optim_uniaxial_stress.png diff --git a/optim_benchmark.py b/examples/optim_benchmark.py similarity index 96% rename from optim_benchmark.py rename to examples/optim_benchmark.py index 993ee4e..48f2655 100644 --- a/optim_benchmark.py +++ b/examples/optim_benchmark.py @@ -63,14 +63,12 @@ jit_updated = init_vals jitted_update = jax.jit(solver.update) for step in tqdm(range(20)): - # jit_updated = solver.update(jit_updated, step + 1) jit_updated = jitted_update(jit_updated, step + 1) true_vel = jit_updated.particles[0].stress def compute_loss(params, *, solver, target_vel, config): - # material = init_simple({"E": params, "density": 1, "id": -1}) material = init_linear_elastic( { "youngs_modulus": params["ym"], @@ -157,10 +155,11 @@ def optax_adam(params, niter, mpm, target_vel, config): ) # ADAM optimizer fig, ax = plt.subplots(1, 2, figsize=(16, 6)) -ax[0].plot(param_list["ym"], "ko", markersize=2, label="E") +ax[0].plot(param_list["ym"], "ko", markersize=2, label="Youngs Modulus") ax[0].grid() ax[0].legend() ax[1].plot(loss_list, "ko", markersize=2, label="Loss") ax[1].grid() ax[1].legend() -plt.show() +# plt.show() +fig.savefig("./examples/optim_uniaxial_stress.png") diff --git a/examples/optim_uniaxial_stress.png b/examples/optim_uniaxial_stress.png new file mode 100644 index 0000000000000000000000000000000000000000..2a6fccd41f9566029743147ade3fd9ca724b5c83 GIT binary patch literal 38565 zcmeFacU+F`|37}&E7@BEg=mPD(jZY%Nh+1fDh-vi(=v+?MbVIkh8AgPqX=p5t&;ZA z-oNK@h1=-EnKQSWzoWiH>PSZhOW`^eKI3@&Jte3gz1Y;7@g-V z)0Nv@FU`LDbZDN>S?TF!f#uEbgHMUSJ>6LJI!(T7<}5Wno+t9LCwvNd{2Ti1Q%+x5 z-l(eY-fbyFZ!In@o!l@y81Px5s=+FSHY1XC3NDQLV!C_J-0a62%45Gpc0b)ys( zQCj;ARh3wLTHG0t#$Vnq>g}jjXJI)uTu^#W2&RZ%AZ!03gXxT~BofERhF9~QJD!+; zy>s+!$A_y;+T$ADC9Kh6&e}hPA`%YIknQvkRk( zPl=57g{z;WX;n)}TX|0C#_4_JyS3aeJ`%jJZ*{ER{e*Cym&uXYCFuv}#fN@{xnUXaW>7Hl5zC6pMExLHih7F?D!+nKjDZ{FC0;a4t?*FiHi)ZNCF@Dw2)#cVH z*j#DZ5*9x;)Z$@n^XAH&p=QODr=>5>@Rsj8al*%DxV#!44BgM#)qOD2OF*$$-8N~g z(_UEI`O@ph;YoeU8X6i+PPQX{v~>LbP5FUaxw+*R?CKOb(lU)wqPT745ypjDg}PHO z+Ql0;`N{~79(jFy+W|qzW2O=6I@K%v0|Gv7yScfShC+E|Q+8guy3~+5%$FW`?6Qd5 zDD%iFufZUBP*Wl=p(DAxwVT0qFxg>X9 z<QW_er9@o#fI;Bl5rC^M+u`t9fGBWb4fDIlE-IMru<^8wUUN_@Cd7C`FsSAzy z6U;gk7#AMW*Vi{M7h9p~-<>5OvuxS2p*Jg%AK$n!_tr|zJr5r~tg-&q-EA`57Ta=* zw<@PG*UuwLMO{5&XlSUj4l_=|`^o3rgJYt4BRD2Tp)$cNwtm*kgkp|pYZ}&!!>X!YLL;4`GcLb&=zer3Ys_pOg~F&{ z+0*EQFW1o2j2ZGwexZBzY(cKyev@~r&TM9WQonWMm8|{31D~f!98R$r^(_>8oH<)0 zrwM}*tXmx3To|I5*7xbtL3w%ktNlB>B<*tQa)dT1KRJ;xt>=jr`+uKqsgW)_NC8d|YB1<*FG`e?1A$i)v#{B@RaMpfCW~;@+A}WiCgz~Q)97}V&~U4{*`}|n*Q_bb$$8w<)1xqh z(YYW@g;yeP7K7E}Wb5JWbjugR0e6nG9epMuCMH&s&dtpY>+l#CRq*0Wb`Tx!!r4KI ziN}^~l0Jf=EV%6w6jWW+D|vHDuNfwLON&Y8=iOc=MM^a;<3-O{CGmV^VK&a1+*fkI zk}VUS31h*3{)nQI_hwPcr z3o|H`g_1rAT5XoQr7v*mi=XZrxp!S!@6!WL41Qr>TX|}EmCdjo-7fP(J9qAM_V@H$ z+gAEw)2>-QJ*7Qc1TXXzbk?MvsyaiV+*!+`Ur};;bymJ1zn8nPt<8}ks(IquNz4)E zj=Yjm@W+n_B=Rm!n{@>LzC3He0q3#a2;1Pu$irBm0sh%J1u#^d3__*7bTXIl0aD)v?wWj9q_=qRK!tH^CtrKeCz zu4^cL%*>Rxv9WP!pttz4sJFRLa(p0NxHjiSA1)6gq#So%!=-^3pPg%I*_0xg`sk>a zn>J(0?33;FZt6u?Xxx@PS`vq4W$*5Dq7M^uTrS~!Fuwqw6T^7GIY7!I0QaRkI@DXk zlRQ4!(GuSho(}_R@^+cxW<1Ai*t-FX)NMVjFS0*wqpAP$r7>l!C56j+$e2wfhV|P0 zsgyep@Hifs*Uo9Hh?5f@sacR(zF;liV}XuDIc8>N{{9N1x{}nGHJFna-IdSVlKQLi z@l*;U)KdnGNs#T}tfC($CKqrJhQ;oR80FFfBYv+)){qY}jTtwv=}BpWQ_9I5f`Y0o zMPd0VW4#=OMMcdI&pgtOj7fjI`*M^<#&XQErql((1z5L*Pv&J*VI^5FU%G3gGj%A- zEZJ|`EZX;xteNiIk`@*g{2htiT>2G!n9afJDK_!L?FlBWCC@u`XDy%)t*xzf=!<;u zV#Anv+d00TI)?G1>SN#B?%%!Zb0YQH2d+$q;s*i2^EeU@VRgq3HHSoN=WoRL(+OD} zb?%8x8I_tbV@7jjLTtQQN4O+=++LNKvpY3ZVRpTQh8j1F)-14{ZBQPi?Lro4tU2zj6 z$;oO0rmZ;$6M8;hUeFBR%vaBuVA}RH!gjQUc&s<&m=ZxM21(H$l4C2*JF3|XYCH@G zhTl+feWDv4x9_5eh;x)tgEV<^Mh#i=8X8ZIu`7DK^NB7BSKXS=)tjL{cGw_gG{&ay zg|~L+=hw2r?3hu`8660V4?Wg+JTwk#R654~kJIFoh$2V=1$eXrc*W+m2 zVx8h}Rd+6JXPbdkhUPc(I0~N1ha8cU^W!}At}&Chx?p^)WBi)BEljeM7r#6%{0@_s z?bsk2e>-_x85}9c76^|#GwrAfa}pjqib|Sugyi@kN;m!q@4_h{=G z;sbaLszM{I>bF`B)Yx-t=P^}UGM6yTws zha15pFmB9QC$FIJF_IO5)rR;`)2+pwB4r6S- zYYPj@-Sj%`?W#wjdfLhrs|Vikl2AtX<2@FHWU=>LT$Xgze3euma=uM3Sk`C#;Wqxdn>+CEmn1uAHRkxl57atxpL$2* zWmaWj7&1$1>(qseS@0!2ScV~?p{!M^Qz&V7t2^9IteiXZ@^%_gn-pQq^!B_bk|9MO z4oFDc)lgbgz9)Yl9S0(S+UKb?hfbY3<>+Kot{2-xrWLtaT_ggl1x!N^95|rdQ-*aX zPb?}~jhcZe&70Zi9P4(jy}dDFr_Gj2*|uZd{O3N~wWo~RSbzJliE04W=)d%-j7E$+<^5AKz9v_QkY6^XkmB$ua;X5*7wQT9q z#u|HuA`D2a{M;P&>5L+4vEGB_0-Iqu!vdteRu^xDwd$%j=1*m~cj@(KmKC|T0~|S$ z`4Z{(y=NHfoJEa|9$g4$YVtJ7x;{W5s|V{as+o^GG#HS22R&T5w?`}Qxz8fj-5%8i zJF<8AjSqT{=bmyBEb{G3>5Gog#z4gPpDBpW*Y1cs7^E>`{9&&wcX~nJOu`(tmCX0r9T1(}4tIYlAFSK$Bj`rkl zVPl&g;XT!uiDY7peC%#p*JU-SyxB%`-_5NB35eIV1q|T~qmBE?<4$epmD$yGWy_W= zJ#RT|WRVs`3H8TibMIU+lYbf>|Lr}BGr4$p-i>~wP(-ZaD)lXr)G|GIWz-^g^}A=- zwy0F+VTjE}`)$W#``RnFMihsUEFEdV6Yr5Pn=G-EqR(X0XFlK*y>O5A|xVX5oSOfK?mQ|*nPE#p4hKuVx)pSa) zO3~VjQccC_4~yGMo{J(}FEp^}_UIh*+wW9#X~w*sbm1{sH#N&(8%(YAMB2wmvqkLq z=Blo`ZzOamJ^+({k+pEXryUSFCK zkC8R0G%0Q>(~T&+$0ElCv#CIL=KY;@Lx?5Asz<)^x^FviiRvUCU1-0{gh)k#B%19n z>u6OmS&0iFGo|V7etPXLPRrirC+(F9-aUhx9+~es_3maa+)yCmTzFALX;)GmIIL>s zS++h;@99ikpj!|mTNb0IlEx=~*v2NgJ)uL0IEMJv$P9uzNO=p5-WwgN>Q4?tl-v9* zgQK`OxsMw;U9d{5eoey#WL5GZ3fhR6$RRyAUMItqElk;tKCHhdUk1}p+;Ig*(bkAW zUs=C>EW*OViTarz5)oxlLoqn9#1BRx3SoJ8T+ue8mLzOMLZ^sp5yA@6x}U2tHE`L`88uO_w+Vzf|-!((}NQhyB>8K-MV#4K|#T_H%D~A)mTK4C5SFwz>TX%1<3$6*|&WpB`B#?h>Fs% zIdTG!an(?Qi7&!Ho<$)1E@5jJ5B!GYVHa!;G?arjxUm%iUdqiU4N-PMlBowN7-1}80&Xh}-u6C+?xZz+y2?Uq^WaZ%#% zlP<(AXFa;G%GKG~x1!>h{E;J`Jtmkea@EP!+`2`qNw6)8vg!EFi6G5Jwlm*A000lS z&4@X6)hdh9mMRO0h)}e)wjMB)lDfsjvT*k1g|nmLlJg&81<<7Cq+MU)XX&|j+0)Y7 z(e@?QG!v=qMI^5Nr_ygr44obN`0m|1zWx9cp0g#7=da=M>>2!8)TRZeOx!pYhxkxv z`-c6ugOHFDSG}{r(=RVCPqSQC-wJ8BhLY>liy|(;hHWDkuBM3))Z6Z5NvwglD~EbG z+@?CoC3=L%yM-62i%fg|Eq6Zy%SxIpJ9g|ic|lV%07kjIRaecjiCJ)Lpl*|DPx8~J zTPogN*Sq#eAos?ytx-r~H~dP!SyRIiI!0Q1zru)d7=S|H6RL)9mmQFiaWZE&-oJGr z#{?^IHcZA_2!RW7wMFG2AtCqg-!D8pj&NsxNAhs^+_`h#Ab#16ArT|)cfbVuA(c_z z>jCU4vu^_%o0PV;wp2`DK)|~%^Ma4kZIT5baGmq<1d3@Io+9a%oxI+JZ*5HC=1s=* z`9(u=hYx=_$LBsu7unLWbg}*D*gld9y)ukP4uX(d3r3Ez5TQUkg)Y-p%|=ZVEw>kf50~J#pKYn5-C~ z=DhT|H9W6!vQ-?y(tvfM*BX%pdT3gTRRmO8CIt&5|RS{%PLvR_4X_WbJ z$A0?-3`JMwEQx|&+-P)a8pUp)Iq!)h<~$&pKzNh%1OYEvvh=ABuMMaRaWe~q21J|* z@b&cSv^5j36vCE*nC|)IaW=tir;aPRVlMW)p2ZM^``S*k;YL!FeVkE)4)W4Wp@#lp znZYvTwLGO?k-DBO-HsI)I?Zwx{O8+kG;}-9?`5zV@Hu?=u(I_}5EN@JN80D@ZP9h& zx<_4~$OsMYh2MWsXdk+e0D!hMxURz1J(YMIT5(v)jYi(X$2_u{1b~>|{WMzypO+yX z2!ds6E{~1CdJa)rGkbc}sdQR0O=A)@Dgb@At+`xKa3nkm6Qq$5h$WeMA=y8_#BRE* zfhB?P8UK%0}EljlN z{?coyG>yROoc-`!j}n}b+1@#FbUK~*LckOR?XIhwkIf|o_Z8KGk)WPe+>W=-GbcDb z5$yEMGU?m!=y`sxiu-`4?WCP1V0uOgAQC@1EE~c|4zNs_Py+z496&Cz{kd)g4Y|ZrFN!N5x%oC#PY~ znO3@GETa|14xZ9y5);Tst*N;#~nF#>=WGwylL#A&O8d`gpYvps7;Ee^rjVG@%iI?Dh@-T?EO4(o@qq{ zmzj9|$$|X%{EqSWFT)2SdF%eO1MH@2x(CdiH&4S|Gcn=E1M-?XW%ixN2lMZx*U3sq zT*t&&_k+ESno2fV33=0`9@@6@-ieiX>|O1Pzh8s#Y9woy0Eg6cZdgh4vB~&dg5>v2 zIoih8&1O-!5CLbj^nUE<&uB9K*k+ne-Qn4~e6ITL!7u0I_%^|ei3Ea}V%ip4o@sv& zho?ztsq&|8!>F-F6lV;w;3NL|0k5c6%jIa&g!efQspgJCaeo_cjQ5QbOB& zG5^FAe%1HLW=I>Mt~LPC)xtK3Hyj=BdwWL^bO9NQUgWsa8LuZ7+T-@GP&K|wyvP}E zYXt)ABA~&J3{?hs{!}jJjZ=PJ-pRu<`%kL!lM^UZP}j(-zM@bik~itXAFLL#dRY{q z9&z&~jb6z!pVCn;=|C#b2PCa1$49ic&A?U#2CYDnBWdY#b!BCDIG{JeNf-KWi|xdh zQqt%c$)B?%7yu7L5|uaQFrN23JVI{VxRIY+70H^^`DUIah?iaTR!7 zQB*a#Pofy$9SNnT%`GkR>go(O$K>QTNX(r+760k6go*9Z+uyzN`X|^&@CD@w@U~P! zaAD?98-IZKTB>WHoNe=S5;_cyOH)&m%9#Kn>i3Tr25`ZB{1v6JxHw}Qjo`QU8{GK} zPYw)9xNK}1YKbVmvre=pOQ1um=FL@iuv~~lPOseO(uZK#=ZT~nPieAVIgiHa(@}nJ zT)|m74RJ1h_$Pex!IN)~zpuv1O`)lV1 zT&MYJN)#y})P5ECH_w|y~19Hht#M}KN~lkwoM@FoHbtgmc} zSNf9c7sYaj$tY1qpw6xNX(f_!!NjoV&t=PF4HTuNr5($)&z`+`G_`IoAq;v5pn>_U zt7uD(ilQ;Iqj>#_Zg#cPVqqR;*|H_?=BgcJ^}s_*zkU1mX;@g{r-wV3CGrDhijmt# zoqRi2!aJ;(6T8f(<3G~d+glc?;Yb!(BO>4XOmu>lMh`N9MY4!}PGC}DGc zpc823;5kb+zHjIjYc&E$s(PU<`s4fehp_s3D?4ly;fEbfEo4vp;t);~>j%4j=ZJn{ z*S;{dBxPhZhhY|wIz}V*>3q`&*q=ajZ-$FOE2vZtxe|cRx=B*~u(o!v4bY$Fm~wp= zCU|Fpf9uDpcn?=Ja3OuUwQ9#cj^u%pFY=F(k!>`um~!bs2nLOSVL&dMG0hFWjQ|>r zGHp}VD~pm4y+)G@*d-8HSkmBE4}{2WOi0j72>z;OUNeJ$6d3sU_V^aEHV6ar`uHK7 zHVjLj3PRY^(&hCu<==OAe>#iaoUw>G)+4*Y8f{Bn0{q^Q(VKNmiO)ClC_xIDZ4*OJ=&#|;l@!6^re1dHb4W!td+ohG3u*# zp7$rHZxXoU8UZmvVGu~itlB!3tRhe%+VvmrKMe`VZ;41@mjJHdQDxOFHQEy>ERR73 zbsaj~k)jTmQBG7;lr^JmAt14}Hq$S^c3}^lgDW7f0Zd6opSAy<2J#QLzA&k|x}y2R zChvj<;c*K9B2D(fBR)fYZ6OHc#hqb2xDXq4-kafk4+gv3AVt+6vFjR0ugWhan5T;z z6^aT9#GS#wBLIv#KpSiQ^wdGzH@=BL`r$8Y=VBOkKDEB_#I?Zt6IJVQXoAs2OebbW z)4D61Tq_|t56jClODuXV)@azc6Nq_(x2+nWZbF;hL!t{-RtNVBI>+p#*Vi(>v6X+* z#>sD#A15Y`+{<2bZJk>NvjmorWWn$f-M|oG!gBtT)I_M?$XZ!IwQm0tnL9Vijf&MHeZ%}jOn8M6mmqFFol$%&CS#$OW z7P~uJXgG8X#FG8x>-*^~_QH-Lp?7Id9@d9eLK~)Qk)EQW;{LUDdy&=wLC}p91gx=9 z{e4<9-}M_4Cia&naE7K;D4>x~$eSaP-isQXP`!`1IqRj{SGJ#izmR+@HKD_|l*JhZ zBB5g%J+sCB-o5wg)%Wf#<>uiz;EWHW5NQ+Apr#=sBZDvq^BIJ8m1*F1C=ze$&tFY- zLJ%I9zp8_stdPLBbc_yhfB|)ojn$;jqFs9kBvclE3o8NYK@ukuCrZ`<;JbJvvIQ{a zQV+L_<(;Y72hfy2`fSQk;z(KYfjGty8!K0LbqIFohzW z^-Dl?z8^^vg?|QNqJjLZRr{oJ*t`-JNK*^gZ9HTb@+vAJJB_|Jg3TuAVT^v|)1Dq< ziNoGP)=2UcDvh!Q6`tCMe2vD&0d+DpVwC=>5 zEzSFC^CQ~i=lBnUTRZ7@TvE%HFaOAuX>K+HrzNk&%csoFx&9!41yHlI^J0r#^m~wN zQ(t^@rY?NEq-)R$^AQH}+KKla{S)yf*?i{?Vt?thzn&#FiR@`!+(Apx+8JBU7E!!4 zrJ5r!pCz`OZsK86cR+m3zbdv|s*^l^@(O&t^eFjjr@$#CclGMk?NU1k8G{FOVaKpU9h#Qu~!hsaQUsol}`=(ch)J+ZDAGXmzl(OpD zx|C{$f5zzzBgTp4uk?u=JjYKmQGoJTuR=dIQ7vMGBtL-_`eqaP6?KQplt9}i`RlpMX&en<`AuC6fRcyYcTo4O~NAR^PJQfHA! zNZhgT-*aikJtUzcubjb&%x`eiW~Y~QTxI*jb90~~fOvjOzRwT4D{m*}2BS6>I3f24 z@eceFsA`M+t0W0SAUq;!p8Bs}br8bjvXdm>mSO|(X^YSVp1q-xJ(onL5)9>jkH)sK%bdTv5pt!>fiJlHPAABR18Nm(r6}Q={)a#kBno81?3EuI7l4kd{uKi$&Dy$f!STXlTec9r7Nk zF%ZdX$@Q~GplZ_7D2pIp6#_5*@saLwEB2VaVl`2JsMlY6gyyVi`_tCZ>DQxVUqMJU>&s7nXBBnabCDf|2d)0*TJ zJN#o2Jgtk?AfUm!RrHu#)4_ClEUSIYpD@#8n3A74AD(=!p2q#=PPMtKF zEN9v`YDzD|G!vH2`AT_Pkvq5UZ0ebb@^1AG&6~*F=&OoKN`eq;Zr!vg@NFAt&}+q= zEB^ptRSF4CW*0*RVJVat#Un@jiAMxCwVGGo6WKd)iVt_3fznCFJABHtnObjeEQ8OB zI7N2hHx5GH8x8r;=$B2xxyWL-NFW>ecc;#*ra_iw3@{A*dDTZ!3Tucfl%urnVh-H% z@(P2f;@S#MVe`(%hYyb#O|#b13nOG^6T|p`JxOaE+2h59N1J3!0M%aWkNY>j-SeQB zc!82d6OKaMUBE2-YnE4es}X4b+OBW)8PQgQ-{fK%{(({b;veJy*uiAZs6VN!fTNkF ze|Nu8j=u;k?rOVFxSAivjKMp5Q|h**G&41je+y{hKs$5HNBpHkIcC()wP%2fl@5f{ zBX@n2iYHQXgnpek+#=X=XW@60s_m<3!W<#;$)N{RA;p|{5m7JW2Z*QMyxH?GU{wb- zoW>@R3gqq*X_L92XKXm+`hb0nV7ED|nB_@5sl9RXpN3{!fv>LzTY-=`544&~y1ngbrN2rjm z=xhw0-OKOTo$<=&9H zyf}l2lM0%?f8#<9#p^6o*8g$x9U&o*GTLT9{25710X!9ftzQ3#%xFODaf>W8a@izP^yPz#hjmsY>#`LqS|GK&pI&#DNpSFs z%STvuFu>e;Wn4tG1Ohphkg?i8>S~HeuAG9Cv=4OU0rp6iNx35*a}(Um!P6t zh1{npFa8kqBg1>v{gSM0%9PoQ58wir#HO$CN%v`hY-mizISQ(XFr^p!pmr_lPwp2) zL^KMWLLtFIVf8f2|A!PeF}0kFMiiWhGgO+cv<4{~#Qesx(tZaN4;#OG93CE?T23No zyR1u;>0Z*Pj9_X)G?NP~iCKckWU>VNkBvfvpU0crb&DrDVR8>xUft#LA1AfsmQD8} zqj$c4pIHK+9r0vusNAHN7n?*)TYALcU02rtJa=`VL^QEf5pxW<%MG^;eSdd#r-`c* zo2ZodhlcjuG1HVmj-02@l_VbNYbIP+6yW5}82Ct{exS;d0b^_i#Yt5jd81SV5Q`#Y zoz6_q+-kwK4*|e^Gknj@E#yr0qeZWOU?+?_4QD~^3q>zM5eK@yE0alh*x+UM=xerj z1b`pil4mCfgGt05q_6~V#a*UeXqSlABpL!a)O;Mm+44je2r@|weN;?dFG1=okV~2n z-58cDo-|Rggh1KgSDLVk+!A?xKoG&%4>^YYim$Fq)y$wiSylGZ(62qQTbl5m&}~2n z{tV*6$8e2X1O)|0+=)y;HQqQEODdGmd3EVd#Zc5fMvNsnbhfL5fa-wU2_;*Qbis@$ zmmnH&)O$X0N|;ebpjMcMO0Fh(4Ro~ju52o%AD^v9%|(b}L=hC7GV$3H!eCd^UV@c~ zls^n$fyo8dLXy39eDq55hbOD~g${8d;jkMvAK&rApf*g(A^qHXKJHD;e()vvO+O9< zmrz16;y`T6SoBK@w5h=d-|GM^Ipc?`U)6YiASqS@nD!zE;ft~a;dB@B61%myzz@Vh~mX)yUtJ>?GLtfIVg8 zJi+>5xfM{M;WBCvSHBzzKL>As@L~yI!6=k*23Ht3cW?0Idi>%eX1BL*z4>pq+ z;?)F6@%VqpF2;#qb*LydnGn`QKsf~K9{TRxJ_2l?Ou;$>wsY0ivcH|qk#M0D2LTcC zr{vh=sRNsiUwF7Qu=_J!CJJe$(E0+CdurBEMevU(tHJYL)_tX!F|w0n40=jz%m3Gm z0jEB$%V6De7rf#29XmF}oYBo*o3A5UwrL~H&;`l@Arf2(41aNr~7 z>JCLqOH1=%WI=Ci=_jE#vi^g2k0-*stpS+`h5KXL8KNH43Y0(4RnD2YLDgAx*VA6DZzRw4Y}iB|a4>D^SDzV~M(58<55+@s z)7^t&Af>+2=a)Kr%&PfTssl;^5s!vYg7ROZdAnid#EV^)m0%N ze*`!k{rh*>;F4uZ6=L-(dr(+nawwZ$__ugdn3`4i-UjC4ae}%^)7<|I^zG)Kn>G9R zhJ%8*D?9UZ36(#E0G^j|YQ~KS_Sbwr*s;9rYX#B6Z~_joyS1&C3E#<=vYSu-4 zHz{OBw)}rHc=#gf8UB!9l+2!RX&Z6FUw?6;3mAlmBU|EWPbNwkF56==B<|UiEK#zb z9_}|2!`Nw`^$YU;Y9`+T+_^}U2@w~hxlNk&O#9~)@7qkfFH+}Vy2WN(us=BOf|-bQR8WTwS#m zQYX9|0+Du+I<|SsS!WwaN=V4+>xbi_l?(Qv{1Utx;Y21XtB_Xe}8wV@DBmxdJq}EN1%IELCrv(1om`845+ZJTer5{ zi%^BMT_OBfFkuxT{-7~fkQ5G23j*I!Ylnb9LTL>2srOu6i%wsFy7OT2&Fh7R?k=+% z9(RQb%l;;T1;|-o(9pH_5zE=)@f=EV?l5m$R0NO1b)?%;&9Gn41N=6paiYmA0XBe^H2K;U;^W7d1S_ z<8fF%bs^I|@RvdOR#Z~d=>4<4Z=9o`2s?}jn~;_e*gH=eVJ>0SF)6qc*@T7FAst06 ze}^rxpi*i-;rg8QQi#+qiabu$*(Dp)rbX}fgs2Ac@(z!HSMhdjuMP9z(Jr(U z_*H$w5`XG=gQ1?|PkpIbt?C?5u>e`TSMe{Uz4;G`P+TDx1>kK$$*=8M=|YGjhj3;cz69IQ=In56%Eo zyQ4YA_EaxG990hKzq zy_erHH7#pmO+Nu#Zy#F-MgnCTA3uD!%S0SlMDoy6a2TkaUiB!TsVUIAh?JvEQ1L(w z9V{Zi!xMD}kH9F&)ya!74#DSthiVX7eo>p3;k#AN{!sNaBLKe7WZkB#o{v&FQgndA z8&zu2D{5g{po+CQdxL8BaoV(H-z%qSIU?Cu65D}g>>vmusgEIID9@cn#si%vq#kVB z4Sz3>w9C=SkPuPedPoZjE{L&uoAURA+snhg{D~;!roKdK*ZqZ91OO0=4oX5=z*}BL zoG)a5F>&7CkFd)EWA^>ulQW`ZBQ?Csgbn(Ymd59H3P=och{)?ldF8t4R4X6H$VMuz6WuNCZVH(?6!Ff zvflw3vh)A9aqKze-^pZhhMe`dFX<1H(GZPGab=#)ztM!k6Kn_$v?Ge%a|iw!FY=|* zA=E(s`9$sBA8axYbZQ9vjO^0LXbu&!7Pbc(z5t zP|?aRrAiUb?ZTz9{mCAP9JWAn7a-vNsB3J2L_HVEI->6)zNB=trMa0%D0LC8p)iwQ zJHhSP`Ri!FvUJP>r4_#NWSZG6A=7=)Z|Uy@-I811mF%)sgL><6aQ_fwY&hxC-_2=Q)`PN?I`@U8li*lh99|g+F&yMdlS2=Lt2GRw2 zS=rq1nEtl%KykPr53A_*HNj zUGYzJO?nRupm=mPo_oWD zUB6f2F;n{`5e7fD+k8Fev)N>QQB+#WqC?;`{HTr)f-@Zm(%=Z2p4F{|G#0gDw};@9 zPw16hq{1zsVjy|#iU5+3qyKaInr{@)InDSW5=A+pM&qlaA<&ewe}`3rz?RLMa{zw# z04&o5xP-AnhAC%tqt$=HO6ZAE?eD*{*ZY35*Q3yGMxfZso*@FlS%R%&4YEKqB{778 zSRW7lc!_bS3{B9VM~n1KLZP%jiVHUdL^Dq-@+I(y3RWLpcK$zOv=wu?_aK~R&4`pP zt&)=wER693dB(B-?cSb~b%7rz}2uI)s_X*>0`oJ^;_7+c!wck|DfOPO(WsVi}Ag zwynrJylsYpnY01VllsmDwgaL`UGLTmjrFHQ;L>zaeNe-LL>k0eAZlK8|d0psEkH5n6PrniZzopK2H;IawT(nm-umt~5_eDF5*scV*+?S(&;%Izyl zET~F&T-h8@!v~!2J$$(8xgg46M=GxObyP=?YWYeAAJc0ztMdGBGLN|E7cWJUk3 zV3@=Uq&zYojkLOi4EDb><A`i}qRP z*Gy~Ca5B{LNC`=`cgbnPZglTJQRV)#^M2%4|DoDhT#c$LMaeb7CX`k!zJ&8WzE0Y5 z@>NZqa?LS-ZoePGEK(OD*p;^AfW-9?`K%8ga?wq1{nsRRh`mwyT+-U3NXCqR!6bo^ z-E^$z2Ld2MCek*RrWaSU0Zx%lVOZRm-oomr7JweS?8O-;qN_qYLz=ZXk1etBME0c& z1dFs|16m>$y679E8&U+lz?$y?Hql5jK-8$N<$O!BJ-aLn>hGFs#wk#=QX5#&ay(#T zS$hB#=AvtO^$RAgUR6Nz2g#vU>DzGQx}e@j!>J8rF|3l##{Zs~ojk%!&aL&UJwIqW z10O+O%lH+=fP6t(w^M6~4jN^uk(%Bq^{y%e9`NkrBPCgA4N_91PXbKw<3Fr^VA_AE zAhM(Vos7V@P+omSj$u6ePcnjRVhpEln&e4Q4}3@6I++mn)w}q#glQD#%adW}dH={ z7n-*;vmb>)uBQWGJh9?12T_pLYxbv*P7_3wgAxKPo|Zd2xu(CTaN>;LOZbXbuU;K> zpj@fZ`?AgTdW1(B$OlE-A|bFUvWwZ<*?2+DiWhXyKbtTkfpQ2NSdV|^wfWy7DK z>QCh;u6|M;m{ZpXGDGEfV{Ml#b$j& z0{>^Ba^8$Azg*2fqhmi368~fb5en()1imP$zWZN=#}W^C49_Er(JczqHX1;NYWL+w zX=4PsiIA&De8G*DZRFYfLgT7Ji<^Utg~>YR{h)_I@~W|J(yj^^DZzi@KVO=sTtHQb zHX=n|xa@}wH)IAbBW_Om_7FQoU~t%~k6BqoAeU1eA26U*APdWQ+{H1%g#hqPbwaKX zhL(^mAaL$gcNoZ0r>-w%y@>tq+^ids|1%{{67H5wh*cX;&R`=wwa70HJo7LQ3w53i9;=w{Ul#1a$!l_bJ;UkO|3%Tt{tH0A|VAz6CA5tT>dHz48-{hPn)5#{o$)fzTmrz)__tr1g z^7QTlhiVI}{FTVfi1(ZufDkEusJ38%}!-MfQaQH%UN!A2-cGD*GI z%AzJTv86P&;^rGp<9s|wU>G;Dg}mr zm)44t)Q)kZv&Omu`COnHoX<7hMV%b^5+Z)5v1gdKg8c}p+nBDlU_$;<^+85yEda8KL!43^u}uuPw{oJn;mhUl ztNp{#w&ONv%*0%)@0N9xgNz&!riaka**!oYp!`BpDscX^+bNDuooy#^U7bB9|u7L@+ zTj`Ll%$Uyb0=3TQG3APtit+(xX2}T&p0+!u(uyk_3~0a6j3`q9T~C*DsD&!3S%^V- zOqyK@4Sc?=**{LY9}@6_bmamLwSv^^k!FRcVWodwjVuifxqLdaNjJr#_b%6YzNBB> zunyPIN1*FXxp&oz>;v)ueybG(dd%%r$&+xhu(Cz^>uYItAV?Sis7#vM0@yANMifJx zLRT<%!3!}&&s>vgPkOC_Aj?6c7x&6>(hU#i>mk+4=q~b{+HDU##){Eb^0LDyp?blU zjy0i}l=b6l!&<{l?6{7_({$}9B0Kgd&NcnwPUn-r;Ts%h6QT6t>%CE!7Px)cIDz_s zu`A2f&`a0ZpXezk@Ur_SR??XsM2$ztd>}WylTokJ&*|4RP+gjpeAH@t z;d*98LEJSdi*@l`&JtJUar)_u%;82qCzLwe|DKr2L3ys%#|`(N*{&`xzszRWzl!$6 zUW|{c1F!w9|FdhJ@{`2GUCrqZvr@N3ovYcCgq|v-vmmLZ=vDTQDNfNXW=Fr_{U~yU z_MzBe(FXLMec5L!!0y~@_V~#^Qei^Zt1BuN*XT_R+l>UQ9FPjVW#c-cxgb^FXh^pP z6P0-+Q9}tv?Kx>VNV?9F#&tcoD;;QjYphW$Z`+hIKAeKSMFFUKC4EFm&ro25H#>zQ zlKR|96S+~y=88zKD?-H;A+K7MaX9zpl8!^-1tH_A?22jkFXhH>^#xrGb6Y?{$8~zK z(behI^A8_AS`yAND4p0MdUzn)hlAl@K8*o#1rK_3S3|)Ri3`Jo6dEBAp&Oc^5!CkG zqb%xBM33vh;q?VvQlvtewz@69R97OPKOi*p;8w-(0(7@+kXFz5P~~yy($q@Ray_r7 zcxYo)j6w*JmydOmfdf}c`VK%egTeq4DssHgP*DOxf|;~y(dBvhHtXX&LnUWznWg3*7Tdzx%x}#&YPw zL=vIm|FuzkrcG(wcYzKV)I22P#?{vSXm+Pg+CZZSkXb@~e5jbTOhhpjlpP#?fqiK9RYp8|w3oqv zoql49A=8H;v{77QQ&Cc)g+;BjK(ObKc{$8!7}UWnD^dooqE%b(R}TYa=TPDC!R#7q zoBo7qC0hGbOW+{@n?q59o<@XQq|HfH19ZBPAVJJ!g@XsTuSDS+>HN@HhmM8wS8jU+j!Eprux)RMdNl4j-a!w% zQxU15aG|e5o8$yYeUwK8QU|rpMm-tJA1yZ?Nm-^R7`2xO%aN{hblI>J+r=<@I&12-P>4j{ z;d4^w&gPW9LREiJjsucO%WihR`S`Rhaya1)D_Ceg^B)5XbIq3yg^=f&+Si z3cfXkbxe6OXvRT0;*^8-iPjiRcX~|GRdudwcl0p7 zCe4`3ia&YqMwD_edkp8*p!^|*jDd^IoQjI)01`YBzmm_AQRSopsgK!`@W_xX|6u>}r|i>y!N%^kaNYUMTZY=TBpI#F#{JMGLo!QKKn z1!ZOTXvhqx4H{u5PQemn0M$jy#uwcg_{j%>#BtL5HUxAXscwK^$ufhu0Q79ofp1C2 z_>va=yA*Uk(XQn?PBhp>q*n&1{U}0ZRhh5QBD=JGx-ACmN7z_dZw?sfiWZyK znSnMtMw(ulMpQ?Vx>y;?C3B=UO!snp?*=L#iUf`eW!NqAWSkvX@{9F-tG1sOcj4KS zjnaaI56_j1x0}+W&X!4E{Ip!USMu&D#^{JK0De~k?~)R)8ATCv5W3V{TfruU-jnEHtoy{@U)-5+Z;&23B(olSxviEDD-D2A zzIJhB=}W^?<_FNuz@j0{tF0N4Z^}jWu9S_CayGWT;<&|tWqO)YUp)a;-WE8#Qc5XU&G@uFma{D^|@HE1ClfsoW#jP2OI z{f&|KtjoJeV`0)X)fm4auy^e#pttV)COT=Wr>EO?=~|mMoEjL!^bliJwj8mYpLz(t ziGW7pd8cnm@5+Moy%%lKxdaq^3==_vywP=Ct1gP=flivF#E`VEnP;m#OzMqNq6D9J zkZ$x@!sAwiNrEe(=^rHY61uG~w7KPfj*Z|$C}}Ea9adFPSS~Ars1f~JUv#J=6kyGm z8+)O4c%4LlwQUN9Dq_pS-n@C^V6--+ zS_T;iWa-ZnN~dcbn8P)r==y{6j693O##}=e{4y$2FG@4>mR3w@5%*p| zqVB38kHd_#FCcDoa;Wn%znNcX9f>@r&%yoSdcU~d1$*-AxMDRg z!2{KEk|sK>NA{=3n#K3$wDdmu#+O;RWOe@6GJczsy6RL+%kE6+`FS}xcg}TWULW4q zlC-mOXhD~i!EJTX4el_7y6c$psIkMropvrYR$M&Fa^qdPww|rk0ZaNj8oQ&kKL2io zrjgE#XcIL@15H&B*yMpZypM{!rWadRz*DjbvYs}YT}cIxk={s7r3%-U76@EL#`A~-x z>v$DK+$!CB>&(lqmwP364RnNOH_S7DBte8GP2l3i(X~6zrx6c-a!{QDPpOzk>n;mUY>&2j_oxp5DpQ5F$QYNHi>rq~GY&dh2Zgy`7_vp?s zZas_9u4vC@Y0sFEEU%}+n}&HA7gc!_73Db>bM+6ojtH!D>i=pS(exgJlfbdKDM?$3 ztOomW972n1K`px2_Ihn%ObA@(dp zU`o`&3Xt?7RMb?p?kz0BJf+Khhj>(`1;3FHNpu>ddb0>k)xy9c?QQLb%TQjj@!&Ir z$)RXZa1=J|ti15xt5s{RAF&M;e(-2u#&KPz6>YD+o|Y>_`{gSsth;RHK8)$BIWk3r zYmbu_U!gI-Pw-V~1=01Rbk9QNWK5R2D+m2Q?R|MP)_c3Rw$ebQK}d!p(XJ>dGxao~ zVvCAGDTU0l&`op67?q46(SRaT=3PP=OCj?VBFT8m)cIW6`#I-4=d82df8Iadwcgfx zR-2yt{!Q2K`@M$Gu(~vG^RAB#+pL#==bY!yaf8krSD`lV7WPOS9oV-- zO_IGTy0t*AM8VD{Shqb#O<#=dPjFCMahS?F*F>IPc1s-QJ1<&zsw}85ux0g};Wo6k ziW&_RU(|3wBdz$^2f3#od}Ej8E&uY>=>4Z}1#(<;y5~5DoXiWP-@W55mPpPlg;*r>=~I^gAdkH`I#iOJ9jKtm(BW}5k?U#^HmJ&T@^JBr&` z&K=5r5@*ZFkG5u;=y2;(v`?)Br=tg=(h@MQAe3FbRhBv&&$25sP9K4Q=$2#gFK~w{ z<)a+SKsBjpV^@n_ayVtpb^LOBJ-Xr7Lo2W4omKA=S<}gNHfnJfhH(MQ7b#IL`FQbH z9PIk?%$Ar3U0Uzvt*={JfL}yf0xi0|0i8)Ay#xLbB5AZoUwIM=&KMerHE^9R6PGW) zW`yBV?Sd;@{0wEMwoEhnCgPXaavkk&3o%dhEE%{z0y@g~R@{G~+-Mi&>UZq%=za?= z#TcZ1HMPID$6u+AmYQMj?!ShSG3T7WkW1UW*R;*xF5g^5u;R~o9n4cK#D}6zjSZ&YOuQ$)2!8i>t>d3!y*Pf8LV#_jhY4?D>5s5bK z!>>%K=Rg0|j9OM<2dk((i3YWt%Fv~UD~DUBIu8LOne$d+OV=P2fbUK(z^9h9*(U8* zG2mEU;&lT*F6Br6P<37&tcODNPjmNaUYt3#uj^9^i_mbBc(qA5!fWKW&wU2>Bo6L_ z{}8oJY!pZ`xG$|*KldmT&JZVzZE-zpSee! z)aUAf)EMm1==M^-;?$8lzuL;>a+~m3NEgoX-xSTbZQvG10$AJaY+mTElaN?6cP^Y7 z&gf56RsWt2vnyn6%a%l-Akci8H=RlCD)y;hdqd1_We~6y_m#BS@o6EKRf4ie+oImA z6dLH!Hm{kQXV9=(jieW0a zX<&)0nwcL!FIQeM=K84oYW)AUZ%8q2^skiLWZs6`oP|}vUo(0nYC}jAJ=@9y8qQ6w zZ9@ri>6ZN=XT!9il+`hxJaK~r;F9?kDcdj$JC*u^XE&E*(eo;oARI|?aQTyxB zeox}YV*|lh+UMgy`vy-J<0>ACCUZ<>)pLJQxmWOXA{Mzp|-o6~kc7tcfE>5k-mao)Bc;k1(X$lq0WG!J6sR^^LWSTdJOZD$b)OUymsqEg|2r+R+K zn@>gqJysk2KSIuIwC;PHEp4-6&#!?^r8HX0?sI~mLkl=4<}a^9W?`?7291Vv`%cs^ zlMUPz&|J~f_Le5sKp<>hhPA_icL74P_n;Lg4PH#%R8;(NJB8(=)o>q- zpRiS(p2I#P(*Ul)WQ2^Em3RU)JzPiR{V==t`&$@HkY9S!m7f{jM33A~dT$I6XAQ8T$=fcd<}WGlbUNU1idFNB>J}TBNyg)*#xnAol>yQ z$_c74>wkYNH>hU82!DQ7o`90;n5EcmC#ez5-GP`FD0VM3M2rZm#oDABr&Vx$cpo3G z-oHCqGV7Cb+CF=G`^rpJiMcB>!bS3e`7SW0O`kM2jbhunI<<&1V_UyZtpVU-KN|Np zX8m#e!qi%Z-WFaZ*A!qEwjd>0#5^6HT+hMV+@=HB0P&M2ESCP^;_`yzN()88-Fp;5s93ze8KEdDb~* zSB|Z#idF2liL=eoo+7{>?gK!4o47QsE%Z*~BLbC_X7S;ASiRo`xYGj}3 zSB6!o>QV6NB;)j?S2A{l}k)xf?e z#wilG#xU&n?#(A~_}=dv!Mc=hM8MGxF{GCS zKZ7|V)(?dvMa;c-`sgU)#%cUaf+TYdu*|%C`Mw{Lp6u@F1 zJ^KA|Mc_kW&%F5fQq|2U4OBN{XF8D{4$aKe@S=7}%unB$ItAH|kMZpk|;2g^_dmXU7o_4=L4D;b_)3T$Z!k&*cbm9$`X=c_U`qf<4(x$}aa2i?By z__*=H#L(1>$Vj1%>E8sk4<0-Plaxz$T^)NZn^tfhbttYu_LMfk5$B-|_l4+^^5TZG3JAl5t z+iTjaQa88tq-S-lgXkUK@89xZpb0W(2!N57`9ggQ?baE!8$_rYEqhY4@WvBtYREtu(%uTzZ{7=hhAn`UhGxtT> zaGCe@^5NYX6IXY+{+^Il_t{M}4Jz`W&J{DRW(K>q7@-G7Ql@!D79cDF$a=`!{0TXn z=1JvFF#DT3c|SBsm9SQ<@ax#~?ShHLwXB2c5BGUo34PJ#R+x@KE=D|XMA1@fH7OoD z#`X<;MoN*W{)TYgP|bn4zUr6_=VDZMnLpHqsc!qScGraC zs|}pB#}uDz+%DYiomJai%E8&sdm#7c+(f^^Pqyz`h6fOzLhdF-)*OtM;W3|9@2p)n z!^15h%Vs{+1hsdd)%HWbuecA4H?zQ?psw@{Vq!%N(4Q#0Yi;Zi6Bids1#_>txw*QX zUGkAUcg8MDVL=}sh?{C_Y8t}9psn$iP3KHcF>dBF8t|f zKgY@|vhDRj%&jd-6G{P6!;Jc6J*9TdeK5p36CTWO6|0IF_i?wJda4x9c{_LReK-1cT9%^%4`FkcRB)c!jXrUG+#4wbBu95~hPP;{Pb{bP{NMgz097m5TdaipvCS%pm?JOG9AB=nL zrYO*8Tg2~f;<#pWoL}CqnPnlUn>DUzLBU2}Vm?gnb zM>a<&ZZ+>;(j1y}91v}4cLCCNG53jTbCdlptx6+C&<|x%eU7H-0ZTt?x1A%mnUdIs zi8^%LE&I8yx1{?vw<{BiR;Uh=(plgT$8EKs@lyhqOw7%*2PuBjR z_WHDZ9@e5)I(m%OiMyu0>KVDcUu$OQ37?~?X`%jX2WnY9IN zjzegu;h4X#yLSaBN=UIU{;{_0>OR9+vf^js%DGppU{n9c;>k9D)gXc=c=iJxs|lJZ zTj^XVPRv9YnD9i!x3<=UTYQ$HKYdNmMsf+>^M9uuLkYhH3>Qyhr)!aWr? zi^EoVXlVtI!lK%HWD*)@j#K+dB9HQHsHwv{OQR2`s2g^`B;u7lG+&_VMi0aF3JUi99;EaAG8?s2YuS#1-FWxpdiu8vt(c$+`x z_KhU6PZ|U?W2A(*6nL@`1u&LYL5I>SH{w!L4*`toZacTI-O~P{`#H4RCr{4X;FQ8F za%#@~KCOsu8L!dF&C=Mg$uB?x+VT|V;VZFbU0cc5FW3|=FO8n(AGJxIuX~L_tH=bG z;r`%}+RDnx&l9udFCvksy3y|#!U*j7m2!?7Sff^HaY#Npa1w7u{&=GT>nUi1ma=#X z$^H`=^fcz80(Sm?E#1?e9^cWV6|Utb^R z9P7!blLN}4ONt7aj>{5~f)!otvidI#tPIu7wxU_M+KoC6PN9g5*(*5Ng7RA#Pj<>UA`Y{p<`ZBR<=g)isb>8Sub1IO)MzM`7dX*K*&=V zX*;6T751Sb413QTBvRUx?E*>B1C^ubqLDl0k6s)mfR)tW?`1h|5nfN0BE4sZ7w9s` zpO}sf!V4Lc_G;B5Z7Z=DCWPKaBUsley}oW5 zbF7b$Hc3Dcqyv>(>LO#7fj+@$%OUdcg^FPnnTUe#V}NH0p39jJ*v0zMW?uDD=0n^ll*10g(!RYjgC*u!s( zh4cUGR_!fGLRtuh4k!DO`l^LM7sx8Pd(Pxf8{f^<@IzP@tDGRq>KMDv8Y;U}kLHQew|LX|(SObP`Nc=pF zyt1n2aO-m}&nx--9vWb4K3Pwar63kiEE2`wCIca1og=x2#LiH%#StnD^GPCyHA^37 zjNB%ldYGU6g+F=+GBE$EwHoSrjLL?Fc&zefKtsx7gCrC5-oz}8dyLWGX5}6HGT?4q zw-pJmYLc{pk?Y*o($E~uv;2CFLv_KA2#^$0!1OQ=ts3FE-eJEnfH{S&iyXMYG}gpS z7e2EB;fgMs??mO#3JdVt+Y!@ozx=Jr8au7HuKP>n9t92^@C;l29WQsm8^qre2!xx$ ziJah4q)$#hLWo2^=c>#GKI`l#%vO6K<3jR!A)%opY6m&jQtK@60*=5+usV<0;~L_= zzJO9(b05eVsn)B`1V9xlTQ30-=sVz9u^+u-FC+7;YTzzB3~XqNlZ>lsu~l7 zJ}j=G2H1_2c2k%#M?z>tmnt#h8^UK;n^=CJ;3ZhnJw=(Ni4@WmWb_o688yN{I%z@c zt<;Zm$s7}xK?d3wpqswBx{&HOkeMsfK&@gbNVo&8^E&!RTzrnv_5gq)zjD{((M?s? z|KRcO;3mu!xw4?SjVNfe&Z-)qTl_SKId$jG3h+#jZmUrtNEWMTO449zVBUP`&oSo? z3X9z?RtdjFw_BrvZnhZR?B5L~S=xjZ6zWUE6>7$F9z1w`>eQ*bfECN8t&>*%IcdN8 zO*T@E>SWvY!eI0+q}+XOB*UgD1R3TveH}eLb0B<;BRr;E%sCb51<4~X>>i}X1L`9pCKd`39iNf+CwLOyJk`Wt z8kmOyy89Zkx<11^eSztBe@5fTD!1`thYtBlM9 z|6;NeW8>g4ENPFzlOyUCWE0i~m7Eq&pBQuF2OkAlvwM3Z+)5w^i}@AoljwCN!GqK? z*f>IJBNZBechluopcjMg@rNfm8}mLBxsqU>H&?5eK_wwA#4HC^ASjBy0A=xkM33OM z{QhWHict%ud9Khm+X|=MNJl4Gw%@ifgxMkfkVG$HhQL>VPw~N(szRvt)lPJ~J$dd1gT8cdtIVR)0sumeFf%f|P(V?Vw zb#;%Sq13CgZTu$C6!D+)#$InE1vJWRh<5{29B2fOFf}enN;MEenuG%29oe?_^q7&P zROaMh$p|@3BBfAMQ~QS}l4!;u7^>f6Q#pV+y_vbWBbFt-lCR3sx!Fc5dq(;Vbg*Rs zF_<@E5@vf9&&#sJs13~`g47T`8>^uC4}Rx{C=@Wdq<0PA0mL1uq)ACTjk6V55*fS>rRyq4u7@bxIEtA)xpT4x@Y3DGr!H_urR;9 z<)vzGRPsI#)4RB^qk#SzIo(BnU3rd7kyD$*93c*kOE4d}!SR3|3e6c0o9Y;KU@2O?sQa49TW#P^KjZEm1Po zW9nI5GZJd(5;#9*cT!`o^0@fmfiXKX<#jg$$CK^pua}=FA)iZrXhrSs5P*|+1m)~@ z0f=gAc4(Mv%s1u?2LSUQ4WMZH3roO&LRw+v za8}l#+y5A2HpY*BRqfPQN{m6k{%eZ{=3ATZF+fKSDh({0E$QM194* z78!X{QcA|5G<|~b#fNW32r4l(G8nz3$1Ql=V=5I`Ke8w%COc}nAq4>dtsEgVRL{To z52cY?grN-zs_Arc68MR*`&~(?(K@>q&IR^e8s^zWz`P+X228xZt`nm(9h||)=}CvK zspKCRIro0lIsOrAK-I6~0x7a(Ve=xOkbicg=bqy0BDRY47|jh(>}@$W_T4F62@B?^ zGQItR0EIVi?1i`*jRh4VCQxQHbrI-pGxq+I$6kLM94*C z=~hmCb*C2zD4L=hQi>(B3{92xhQUO8)F;8{!Gxdz6+&f+G1g*Gkd-@;3CI-k*wkX0mn{YSyA(aw zb}p@J*xc3C)nn0Y9`0{${#u`2(a&gwy%w$wt`=6ryMJ!otgL>OhTtCrrE5l}O34z-abomtN_Si2&QDNPMOt#*{vmwbAtVSm2mYe_sAStrwUU|`W zw|j1dDL->P5aWTUnM5_hOc7)z(3Z2HBRoGl(02Cn>`*%`Wit-izt$Z(pvWlPL;C0?jhUeAx?%j;auoI8T*EukHhgPu%MWSTkUF1KZNk&B> zxL2Va9jAuCT&W?rRcvCSMDD`=9Vl5h0n{#qm!dO!s9cJEhHKWdY8h$i5E6aqRA~O; z`vt}(ByVSoXWw4imdLEhoBpV}83o_WqakROODSrSdLE0HU z*GwqSp^43RitBFVMMeF&KTP&r7`$_00j4*GpGqgzghk_*?uYx$dP*X|qL|!$sIF#H zns!iTtD9($5=~}KLE)jGqq8#v2{lGOf-}L{+1bgX!kT5j_kX$5C8xb5`GY0@S*ZAT zXUv@;SyJur--?L8D~kW)s^kBE;IBJC1>oN|