89 lines
3.9 KiB
Python
89 lines
3.9 KiB
Python
|
# -*- coding: utf-8 -*-
|
||
|
|
'''
|
||
|
|
@ Copyright (c) 2022 by Zeyu Zhang, All Rights Reserved.
|
||
|
|
@ Author : Zeyu Zhang
|
||
|
|
@ Email : zhangzeyu_work@outlook.com
|
||
|
|
@ Date : 2022-10-25 09:30:27
|
||
|
|
@ LastEditTime : 2022-10-25 09:38:08
|
||
|
|
@ FilePath : /ZZY_CODE/Env_JAX/IDRL/Linear_TO/TO_FEA.py
|
||
|
|
@
|
||
|
|
@ Description :
|
||
|
|
@ Reference :
|
||
|
|
'''
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
from numpy import float64
|
||
|
|
import jax.numpy as jnp
|
||
|
|
import jax.ops as jops
|
||
|
|
import jax.scipy.linalg as jsl
|
||
|
|
from jax import jit
|
||
|
|
from jax.config import config
|
||
|
|
config.update("jax_enable_x64", True)
|
||
|
|
|
||
|
|
|
||
|
|
def linear_stiffness_matrix(young, poisson, ele_thickness):
|
||
|
|
young, poisson, h = young, poisson, ele_thickness
|
||
|
|
k = np.array([1/2-poisson/6, 1/8+poisson/8, -1/4-poisson/12, -1/8+3*poisson/8,
|
||
|
|
-1/4+poisson/12, -1/8-poisson/8, poisson/6, 1/8-3*poisson/8])
|
||
|
|
stiffness_ele = h/(1-poisson**2)*np.array([[k[0], k[1], k[2], k[3], k[4], k[5], k[6], k[7]],
|
||
|
|
[k[1], k[0], k[7], k[6],
|
||
|
|
k[5], k[4], k[3], k[2]],
|
||
|
|
[k[2], k[7], k[0], k[5],
|
||
|
|
k[6], k[3], k[4], k[1]],
|
||
|
|
[k[3], k[6], k[5], k[0],
|
||
|
|
k[7], k[2], k[1], k[4]],
|
||
|
|
[k[4], k[5], k[6], k[7],
|
||
|
|
k[0], k[1], k[2], k[3]],
|
||
|
|
[k[5], k[4], k[3], k[2],
|
||
|
|
k[1], k[0], k[7], k[6]],
|
||
|
|
[k[6], k[3], k[4], k[1],
|
||
|
|
k[2], k[7], k[0], k[5]],
|
||
|
|
[k[7], k[2], k[1], k[4],
|
||
|
|
k[3], k[6], k[5], k[0]]
|
||
|
|
], dtype=float64)
|
||
|
|
return stiffness_ele
|
||
|
|
|
||
|
|
|
||
|
|
def density_matrix(rou, ele_length, ele_width, ele_thickness):
|
||
|
|
ele_volume = ele_length*ele_width*ele_thickness
|
||
|
|
density_ele = ele_volume / 36 * np.array([[4, 0, 2, 0, 1, 0, 2, 0],
|
||
|
|
[0, 4, 0, 2,
|
||
|
|
0, 1, 0, 2],
|
||
|
|
[2, 0, 4, 0,
|
||
|
|
2, 0, 1, 0],
|
||
|
|
[0, 2, 0, 4,
|
||
|
|
0, 2, 0, 1],
|
||
|
|
[1, 0, 2, 0,
|
||
|
|
4, 0, 2, 0],
|
||
|
|
[0, 1, 0, 2,
|
||
|
|
0, 4, 0, 2],
|
||
|
|
[2, 0, 1, 0,
|
||
|
|
2, 0, 4, 0],
|
||
|
|
[0, 2, 0, 1,
|
||
|
|
0, 2, 0, 4],
|
||
|
|
], dtype=float64)
|
||
|
|
return density_ele
|
||
|
|
|
||
|
|
|
||
|
|
# @jit
|
||
|
|
def FEA(young, young_min, ke, xPhys_1D, freedofs, penal, idx, fext, s_K_predefined, dispTD_predefined):
|
||
|
|
|
||
|
|
def kuf_solve(fext_1D_free, s_K_free):
|
||
|
|
u_free = jsl.solve(s_K_free, fext_1D_free,
|
||
|
|
sym_pos=True, check_finite=False)
|
||
|
|
return u_free
|
||
|
|
|
||
|
|
def young_modulus(xPhys_1D, young, young_min, ke, penal):
|
||
|
|
|
||
|
|
sK = ((ke.flatten()[jnp.newaxis]).T * (young_min + (xPhys_1D) ** penal * (young-young_min))
|
||
|
|
).flatten(order='F')
|
||
|
|
return sK
|
||
|
|
|
||
|
|
sK = young_modulus(xPhys_1D, young, young_min, ke, penal)
|
||
|
|
s_K = jops.index_add(s_K_predefined, idx, sK)
|
||
|
|
s_K_free = s_K[freedofs, :][:, freedofs]
|
||
|
|
fext_free = fext[freedofs]
|
||
|
|
u_free = kuf_solve(fext_free, s_K_free)
|
||
|
|
dispTD = jops.index_add(dispTD_predefined, freedofs, u_free)
|
||
|
|
return dispTD
|