del_candle/
diffcoord_polyloop2.rs

1use std::ops::Deref;
2
3use candle_core::{CpuStorage, Layout, Shape, Tensor};
4
5pub struct Layer {}
6
7impl candle_core::CustomOp1 for crate::diffcoord_polyloop2::Layer {
8    fn name(&self) -> &'static str {
9        "polyloop to diffcoord"
10    }
11
12    fn cpu_fwd(
13        &self,
14        storage: &CpuStorage,
15        layout: &Layout,
16    ) -> candle_core::Result<(CpuStorage, Shape)> {
17        let (num_vtx, two) = layout.shape().dims2()?;
18        assert_eq!(two, 2);
19        let vtx2xy = storage.as_slice::<f32>()?;
20        let mut vtx2diff = vec![0f32; num_vtx * 2];
21        for i_edge in 0..num_vtx {
22            let i0_vtx = (i_edge + num_vtx - 1) % num_vtx;
23            let i1_vtx = i_edge;
24            let i2_vtx = (i_edge + 1) % num_vtx;
25            for i_dim in 0..2 {
26                vtx2diff[i1_vtx * 2 + i_dim] -= 0.5f32 * vtx2xy[i0_vtx * 2 + i_dim];
27                vtx2diff[i1_vtx * 2 + i_dim] += 1.0f32 * vtx2xy[i1_vtx * 2 + i_dim];
28                vtx2diff[i1_vtx * 2 + i_dim] -= 0.5f32 * vtx2xy[i2_vtx * 2 + i_dim];
29            }
30        }
31        let shape = candle_core::Shape::from((num_vtx, 2));
32        let storage = candle_core::WithDType::to_cpu_storage_owned(vtx2diff);
33        Ok((storage, shape))
34    }
35
36    /// This function takes as argument the argument `arg` used in the forward pass, the result
37    /// produced by the forward operation `res` and the gradient of the result `grad_res`.
38    /// The function should return the gradient of the argument.
39    fn bwd(
40        &self,
41        _vtx2xy: &Tensor,
42        _vtx2diff: &Tensor,
43        dw_vtx2diff: &Tensor,
44    ) -> candle_core::Result<Option<Tensor>> {
45        let (num_vtx, two) = dw_vtx2diff.shape().dims2()?;
46        assert_eq!(two, 2);
47        let dw_vtx2diff = dw_vtx2diff.storage_and_layout().0;
48        let dw_vtx2diff = match dw_vtx2diff.deref() {
49            candle_core::Storage::Cpu(cpu_tri2vtx) => cpu_tri2vtx.as_slice::<f32>()?,
50            _ => panic!(),
51        };
52        let mut dw_vtx2xy = vec![0f32; num_vtx * 2];
53        for i_edge in 0..num_vtx {
54            let i0_vtx = (i_edge + num_vtx - 1) % num_vtx;
55            let i1_vtx = i_edge;
56            let i2_vtx = (i_edge + 1) % num_vtx;
57            for i_dim in 0..2 {
58                dw_vtx2xy[i0_vtx * 2 + i_dim] -= 0.5 * dw_vtx2diff[i1_vtx * 2 + i_dim];
59                dw_vtx2xy[i1_vtx * 2 + i_dim] += 1.0 * dw_vtx2diff[i1_vtx * 2 + i_dim];
60                dw_vtx2xy[i2_vtx * 2 + i_dim] -= 0.5 * dw_vtx2diff[i1_vtx * 2 + i_dim];
61            }
62        }
63        let dw_vtx2xy = Tensor::from_vec(
64            dw_vtx2xy,
65            candle_core::Shape::from((num_vtx, 2)),
66            &candle_core::Device::Cpu,
67        )?;
68        Ok(Some(dw_vtx2xy))
69    }
70}
71
72#[test]
73fn edge_length_constraint() -> anyhow::Result<()> {
74    let num_vtx = 16;
75    // let edge_length = 2.0f32 * std::f32::consts::PI / num_vtx as f32;
76    let mut vtx2xy = del_msh_core::polyloop2::from_circle(1.0, num_vtx);
77    {
78        use rand::Rng;
79        let mut rng = rand::thread_rng();
80        for mut vtx in vtx2xy.column_iter_mut() {
81            vtx += nalgebra::Vector2::<f32>::new(rng.gen(), rng.gen());
82        }
83    }
84    let vtx2xy = {
85        candle_core::Var::from_slice(
86            vtx2xy.as_slice(),
87            candle_core::Shape::from((vtx2xy.ncols(), 2)),
88            &candle_core::Device::Cpu,
89        )
90        .unwrap()
91    };
92    for iter in 0..200 {
93        let render = Layer {};
94        let vtx2diff = vtx2xy.apply_op1(render)?;
95        assert_eq!(vtx2diff.shape(), vtx2xy.shape());
96        let loss_straight = vtx2diff.sqr()?.sum_all()?;
97        {
98            let val_loss = loss_straight.to_vec0::<f32>()?;
99            dbg!(val_loss);
100        }
101        let grad = loss_straight.backward()?;
102        let dw_vtx2xyz = grad.get(&vtx2xy).unwrap();
103        let _ = vtx2xy.set(&vtx2xy.as_tensor().sub(&(dw_vtx2xyz * 0.1)?)?);
104        if iter % 10 == 0 {
105            let vtx2xy: Vec<_> = vtx2xy.flatten_all()?.to_vec1::<f32>()?;
106            del_msh_core::io_obj::save_vtx2xyz_as_polyloop(
107                format!("target/polyloop_{}.obj", iter),
108                &vtx2xy,
109                2,
110            )?;
111        }
112    }
113    Ok(())
114}
del_candle/diffcoord_polyloop2.rs

del_candle/
diffcoord_polyloop2.rs
