simulate.py

import argparse
import os

import torch
from plyfile import PlyData
import numpy as np
from mmengine.config import Config
from mmengine import fileio

from mpm.geometric_model import GeometricModel


device = torch.device("cuda:0")
torch.manual_seed(123)
np.random.seed(123)


def config_parser():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument(
        "--config",
        type=str,
        required=True,
        help="config file path defining simulation parameters, including material type and collider",
    )
    parser.add_argument(
        "--particles",
        type=str,
        default="assets/particles/cross.ply",
        help="ply file path defining initial particle positions",
    )
    return parser


if __name__ == "__main__":
    parser = config_parser()
    args = parser.parse_args()

    cfg = Config.fromfile(args.config)["cfg"]
    assert cfg["write_out"]

    if cfg["dtype"] == "float32":
        dtype = torch.float32
    else:
        dtype = torch.float64

    # load initial particle positions
    plydata = PlyData.read(args.particles)
    vertices = plydata["vertex"]
    positions = np.vstack([vertices["x"], vertices["y"], vertices["z"]]).T
    init_particles_pos = torch.tensor(positions, dtype=dtype)

    model = GeometricModel(**cfg, device=device)
    model.to(device)

    out = torch.tensor(torch.nan)
    dt = cfg["dt"]
    while out is not None and out.isnan():
        model.dynamic_observer.simulator.set_dt(dt)
        out = model.forward(cfg["n_frames"], [init_particles_pos])
        if out is not None and out.isnan():
            print("CFL condition not satisfied. Reducing dt and retrying.")
            dt = dt / 2

    simulation_output= model.dynamic_observer.base_dir
    print("Simulation data saved at " + simulation_output)
    cfg["particles"] = args.particles
    fileio.dump(cfg, os.path.join(simulation_output, "cfg.json"))