utils/local_utils.py

import torch

def apply_grid_bc_w_freeze_pts(grid_size, grid_lim, freeze_pts, mpm_solver):

    device = freeze_pts.device

    grid_pts_cnt = torch.zeros(
        (grid_size, grid_size, grid_size), dtype=torch.int32, device=device
    )

    dx = grid_lim / grid_size
    inv_dx = 1.0 / dx

    freeze_pts = (freeze_pts * inv_dx).long()

    for x, y, z in freeze_pts:
        grid_pts_cnt[x, y, z] += 1

    freeze_grid_mask = grid_pts_cnt >= 1

    freeze_grid_mask_int = freeze_grid_mask.type(torch.int32)

    number_freeze_grid = freeze_grid_mask_int.sum().item()
    print("number of freeze grid", number_freeze_grid)

    mpm_solver.enforce_grid_velocity_by_mask(freeze_grid_mask_int)

    # add debug section:

    return freeze_grid_mask