diff --git a/diffmpm/solver.py b/diffmpm/solver.py index 3b1ae01..197bad5 100644 --- a/diffmpm/solver.py +++ b/diffmpm/solver.py @@ -186,6 +186,7 @@ def _step(i, data): def _write(self, i): arrays = {} + arrays["nodal_position"]=self.mesh.elements.nodes.loc for name in self.__particle_props: arrays[name] = jnp.array( [ diff --git a/diffmpm/writers.py b/diffmpm/writers.py index fdc5cd2..6617d94 100644 --- a/diffmpm/writers.py +++ b/diffmpm/writers.py @@ -1,6 +1,7 @@ import abc import logging from pathlib import Path +from pyevtk.hl import pointsToVTK, gridToVTK from typing import Tuple, Annotated, Any from jax.typing import ArrayLike @@ -52,3 +53,77 @@ def write( filepath.parent.mkdir(parents=True) np.savez(filepath, **arrays) logger.info(f"Saved particle data for step {step} at {filepath}") + + +class VTKWriter(_Writer): + def write(self, args, transforms, **kwargs): + arrays, step = args + max_digits = int(np.log10(kwargs["max_steps"])) + 1 + if step == 0: + req_zeros = max_digits - 1 + else: + req_zeros = max_digits - (int(np.log10(step)) + 1) + fileno = f"{'0' * req_zeros}{step}" + filepath = Path(kwargs["out_dir"]).joinpath(f"particles_{fileno}") + if not filepath.parent.is_dir(): + filepath.parent.mkdir(parents=True) + coords = np.array(arrays["loc"]) + x = coords[:, :, 0].flatten() + y = coords[:, :, 1].flatten() + z = np.zeros_like(x) + strain = np.array(arrays["strain"]) + stress = np.array(arrays["stress"]) + velocity = np.array(arrays["velocity"]) + velocity_x = velocity[:, :, 0].flatten() + veloctiy_y = velocity[:, :, 1].flatten() + velocity_z = np.zeros_like(velocity_x) + strain_xx = strain[:, :, 0].flatten() + strain_yy = strain[:, :, 1].flatten() + strain_zz = strain[:, :, 2].flatten() + strain_xy = strain[:, :, 3].flatten() + strain_xz = strain[:, :, 4].flatten() + strain_yz = strain[:, :, 5].flatten() + stress_xx = stress[:, :, 0].flatten() + stress_yy = stress[:, :, 1].flatten() + stress_zz = stress[:, :, 2].flatten() + stress_xy = stress[:, :, 3].flatten() + stress_xz = stress[:, :, 4].flatten() + stress_yz = stress[:, :, 5].flatten() + x = coords[:, :, 0].flatten() + y = coords[:, :, 1].flatten() + z = np.zeros_like(x) + pointsToVTK( + f"{filepath}", + np.array(x), + np.array(y), + z, + data={ + "strain_xx": strain_xx, + "strain_yy": strain_yy, + "strain_zz": strain_zz, + "strain_xy": strain_xy, + "strain_xz": strain_xz, + "strain_yz": strain_yz, + "stress_xx": stress_xx, + "stress_yy": stress_yy, + "stress_zz": stress_zz, + "stress_xy": stress_xy, + "stress_xz": stress_xz, + "stress_yz": stress_yz, + "velocity_x": velocity_x, + "velocity_y": veloctiy_y, + "velocity_z": velocity_z, + }, + ) + # if "nodal_postions" in arrays: + if(step==0): + x_nodal = arrays["nodal_position"][:, :, 0].flatten() + y_nodal = arrays["nodal_position"][:, :, 1].flatten() + z_nodal = np.zeros_like(x_nodal) + gridToVTK( + f"{filepath}_node_loc", + np.array(x_nodal), + np.array(y_nodal), + z_nodal, + ) + logger.info(f"Saved particle data for step {step} at {filepath}") diff --git a/docs/assets/uniaxial_nodal_force_with_mesh.png b/docs/assets/uniaxial_nodal_force_with_mesh.png new file mode 100644 index 0000000..4fd18ff Binary files /dev/null and b/docs/assets/uniaxial_nodal_force_with_mesh.png differ diff --git a/docs/assets/uniaxial_nodal_forces.jpeg b/docs/assets/uniaxial_nodal_forces.jpeg new file mode 100644 index 0000000..d0b350a Binary files /dev/null and b/docs/assets/uniaxial_nodal_forces.jpeg differ diff --git a/docs/assets/uniaxial_stress.jpeg b/docs/assets/uniaxial_stress.jpeg new file mode 100644 index 0000000..53bfacf Binary files /dev/null and b/docs/assets/uniaxial_stress.jpeg differ diff --git a/docs/assets/uniaxial_stress_with_mesh.png b/docs/assets/uniaxial_stress_with_mesh.png new file mode 100644 index 0000000..7455c78 Binary files /dev/null and b/docs/assets/uniaxial_stress_with_mesh.png differ diff --git a/docs/assets/uniaxial_traction.jpeg b/docs/assets/uniaxial_traction.jpeg new file mode 100644 index 0000000..469288a Binary files /dev/null and b/docs/assets/uniaxial_traction.jpeg differ diff --git a/docs/assets/uniaxial_traction_with_mesh.png b/docs/assets/uniaxial_traction_with_mesh.png new file mode 100644 index 0000000..38ee207 Binary files /dev/null and b/docs/assets/uniaxial_traction_with_mesh.png differ