Pytorch autograd for 2D linear elasticity

Daniel_Chou0916: <p>Suppose I have the displacement matrix containing n point-wise displacement<br>
which comes from the coordinates<br>
I’m trying to solve the linear elastic problem, with the equilibrium governing equations given as follows:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/5/95033ba382821f63da00a85135aac7209728a499.png" data-download-href="https://discuss.pytorch.org/uploads/default/95033ba382821f63da00a85135aac7209728a499" title="CodeCogsEqn (6)"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/9/5/95033ba382821f63da00a85135aac7209728a499_2_324x250.png" alt="CodeCogsEqn (6)" data-base62-sha1="lge4jT0EBv5fnWFxSmquvms1Qtb" width="324" height="250" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/9/5/95033ba382821f63da00a85135aac7209728a499_2_324x250.png, https://discuss.pytorch.org/uploads/default/optimized/3X/9/5/95033ba382821f63da00a85135aac7209728a499_2_486x375.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/9/5/95033ba382821f63da00a85135aac7209728a499_2_648x500.png 2x" data-dominant-color="000000"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">CodeCogsEqn (6)</span><span class="informations">828×638 17.2 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>However, the Pytorch neural networks leads to very bad results, I was still suspecting I made some mistakes on calculating the gradient by torch.autograd function after checking my codes several times. Could anyone please give me some suggestions?</p>
<ol>
<li>For dislacement to strain, I have two functions:</li>
</ol>
<pre><code class="lang-auto">def calculate_Jacobians(u,X):
    J = []
    _,d=u.shape
    for i in range(d):
        u_i = u[:, i]
        grad = torch.autograd.grad(u_i, X,grad_outputs=torch.ones_like(u_i).to(u.device), create_graph=True)[0]
        #print('grad:',grad)
        J.append(grad)
    return torch.stack(J, dim=1)

def u_to_e(u,X):
    Jacobian=calculate_Jacobians(u,X)
    strain=0.5*(Jacobian+torch.transpose(Jacobian,-1,-2))
    return strain
</code></pre>
<p>The Jacobian is presented in a batch manner which has shape (n,2,2).<br>
Hence, the strain is also a batch tensor (n,2,2)</p>
<ol start="2">
<li>For strain to stress, I have three functions:</li>
</ol>
<pre><code class="lang-auto">def batch_eye(x):
    n,d=x.shape[:-1]
    return torch.eye(d, d, dtype=x.dtype, device=x.device).unsqueeze(0).repeat(n, 1, 1)
def batch_trace(x):
    return x.diagonal(offset=0, dim1=-1, dim2=-2).sum(-1,keepdim=True).unsqueeze(-1)

def e_to_sigma(strain,lamb,mu):
    return lamb*batch_eye(strain)*batch_trace(strain)+2*mu*strain
</code></pre>
<p>And therefore, the stress is a batch tensor (n,2,2)</p>
<ol start="3">
<li>the stress divergence<br>
After obtaining the stress, I have to calculate the divergence of stress wrt to all directions<br>
I have one function for this:</li>
</ol>
<pre><code class="lang-auto">def point_wise_divergence(stress,X):
    _,d=X.shape
    div=[]
    for i in range(d):
        sigma_col=stress[:,:,i]
        grad=torch.autograd.grad(sigma_col, X,grad_outputs=torch.ones_like(sigma_col), create_graph=True)[0]
        div.append(grad)
    Div_comp=torch.stack(div,dim=1)
    Div=Div_comp.sum(-1)
    return Div
</code></pre>
<p>For the displacement gradient, it’s easy to have an example problem from text books to verify my code, but I’m not sure I did it correctly for stress divergence.</p>