mnabian commented 4 months ago

Modulus Pull Request

Description

Add a warp based SDF calculation utility.

Checklist

[x] I am familiar with the Contributing Guidelines.
[x] New or existing tests cover these changes.
[x] The documentation is up to date with these changes.
[x] The CHANGELOG.md is up to date with these changes.
[ ] An issue is linked to this pull request.

Dependencies

mnabian commented 4 months ago

/blossom-ci

mnabian commented 4 months ago

/blossom-ci

loliverhennigh commented 4 months ago

Probably not needed for this PR but I have found the following structure to be extremely helpful for keeping track of warp kernels and functions. The bellow code is to update the electric field in a FDTD solver. You can construct the function with electric_field_update = ElectricFeildUpdate(), and then call with electric_field_update(...). If you have this kind of structure all the warp kernels and function can be contained within. If you have kinda meta programming and want to change things inside the warp kernel then this can work as well. For example, suppose we wanted to make it work with 3D and 2D. You could implement the 2D kernels in the class and then when you make the function/operator electric_field_update = ElectricFieldUpdate(dim=3). Anyways, maybe there is a common structure we can use for this going forward.


class ElectricFieldUpdate(Operator):
    """
    Electric field update operator
    """

    @wp.func
    def _sample_electric_property(
        solid_id: wp.array4d(dtype=wp.uint8),
        material_property: wp.array(dtype=wp.float32),
        i: wp.int32,
        j: wp.int32,
        k: wp.int32,
    ):
        # Get material property
        prop_0_0_1 = material_property[wp.int32(solid_id[0, i - 1, j - 1, k])]
        prop_0_1_0 = material_property[wp.int32(solid_id[0, i - 1, j, k - 1])]
        prop_0_1_1 = material_property[wp.int32(solid_id[0, i - 1, j, k])]
        prop_1_0_0 = material_property[wp.int32(solid_id[0, i, j - 1, k - 1])]
        prop_1_0_1 = material_property[wp.int32(solid_id[0, i, j - 1, k])]
        prop_1_1_0 = material_property[wp.int32(solid_id[0, i, j, k - 1])]
        prop_1_1_1 = material_property[wp.int32(solid_id[0, i, j, k])]

        # Get average property
        prop_x = (prop_1_1_1 + prop_1_1_0 + prop_1_0_1 + prop_1_0_0) / 4.0
        prop_y = (prop_1_1_1 + prop_1_1_0 + prop_0_1_1 + prop_0_1_0) / 4.0
        prop_z = (prop_1_1_1 + prop_1_0_1 + prop_0_1_1 + prop_0_0_1) / 4.0

        return wp.vec3(prop_x, prop_y, prop_z)

    @wp.kernel
    def _update_electric_field(
        electric_field: wp.array4d(dtype=wp.float32),
        magnetic_field: wp.array4d(dtype=wp.float32),
        impressed_current: wp.array4d(dtype=wp.float32),
        solid_id: wp.array4d(dtype=wp.uint8),
        eps_mapping: wp.array(dtype=wp.float32),
        sigma_e_mapping: wp.array(dtype=wp.float32),
        spacing: wp.vec3f,
        dt: wp.float32,
        nr_ghost_cells: wp.int32,
    ):
        # get index
        i, j, k = wp.tid()

        # Skip ghost cells
        i += nr_ghost_cells
        j += nr_ghost_cells
        k += nr_ghost_cells

        # get properties
        eps = ElectricFieldUpdate._sample_electric_property(
            solid_id, eps_mapping, i, j, k
        )
        sigma_e = ElectricFieldUpdate._sample_electric_property(
            solid_id, sigma_e_mapping, i, j, k
        )

        # Get coefficients
        _denom = 2.0 * eps + sigma_e * dt
        c_ee = wp.cw_div(2.0 * eps - sigma_e * dt, _denom)
        c_eh = (2.0 * dt) / wp.cw_mul(spacing, _denom)
        c_ej = (-2.0 * dt) / _denom

        # Get curl of magnetic field
        curl_h_x = (magnetic_field[2, i, j, k] - magnetic_field[2, i, j - 1, k]) - (
            magnetic_field[1, i, j, k] - magnetic_field[1, i, j, k - 1]
        )
        curl_h_y = (magnetic_field[0, i, j, k] - magnetic_field[0, i, j, k - 1]) - (
            magnetic_field[2, i, j, k] - magnetic_field[2, i - 1, j, k]
        )
        curl_h_z = (magnetic_field[1, i, j, k] - magnetic_field[1, i - 1, j, k]) - (
            magnetic_field[0, i, j, k] - magnetic_field[0, i, j - 1, k]
        )
        curl_h = wp.vec3(curl_h_x, curl_h_y, curl_h_z)

        # compute new electric field
        e = wp.vec3f(
            electric_field[0, i, j, k],
            electric_field[1, i, j, k],
            electric_field[2, i, j, k],
        )
        cur = wp.vec3f(
            impressed_current[0, i, j, k],
            impressed_current[1, i, j, k],
            impressed_current[2, i, j, k],
        )
        new_e = wp.cw_mul(c_ee, e) + wp.cw_mul(c_eh, curl_h) + wp.cw_mul(c_ej, cur)

        # Set electric field
        electric_field[0, i, j, k] = new_e[0]
        electric_field[1, i, j, k] = new_e[1]
        electric_field[2, i, j, k] = new_e[2]

    def __call__(
        self,
        electric_field: wp.array4d(dtype=wp.float32),
        magnetic_field: wp.array4d(dtype=wp.float32),
        impressed_current: wp.array4d(dtype=wp.float32),
        solid_id: wp.array4d(dtype=wp.uint8),
        eps_mapping: wp.array(dtype=wp.float32),
        sigma_e_mapping: wp.array(dtype=wp.float32),
        spacing: Union[float, tuple[float, float, float]],
        dt: float,
        nr_ghost_cells: int = 1,
    ):
        # Launch kernel
        wp.launch(
            self._update_electric_field,
            inputs=[
                electric_field,
                magnetic_field,
                impressed_current,                                                                                                                                                                                                                                                                                                            
                solid_id,                                                                                                                                                                                                                                                                                                                     
                eps_mapping,                                                                                                                                                                                                                                                                                                                  
                sigma_e_mapping,                                                                                                                                                                                                                                                                                                              
                spacing,                                                                                                                                                                                                                                                                                                                      
                dt,                                                                                                                                                                                                                                                                                                                           
                nr_ghost_cells,                                                                                                                                                                                                                                                                                                               
            ],                                                                                                                                                                                                                                                                                                                                
            dim=[x - 2 * nr_ghost_cells for x in solid_id.shape[1:]],                                                                                                                                                                                                                                                                         
        )                                                                                                                                                                                                                                                                                                                                     

        return electric_field