EMsFEA-net/topopt_EMsFEA.py

from __future__ import division
import numpy as np
from scipy.sparse import coo_matrix
from scipy.sparse.linalg import spsolve
from matplotlib import colors
import matplotlib.pyplot as plt


import torch
import torch.nn as nn
import torch.nn.functional as F

from utils.data_standardizer import standardization
from utils.data_loader import data_loader


def top_EMsFEA(nelx,nely,volfrac,penal,rmin,ft,mod_idx,m):
	print("Minimum compliance problem with OC")
	print("ndes: " + str(nelx) + " x " + str(nely))
	print("volfrac: " + str(volfrac) + ", rmin: " + str(rmin) + ", penal: " + str(penal))
	print("Filter method: " + ["Sensitivity based","Density based"][ft])
	# Max and min stiffness
	Emin=1e-9
	Emax=1.0
	c_nelx=int(nelx/m)
	c_nely=int(nely/m)
	# dofs:
	ndof = 2*(nelx+1)*(nely+1)
	coarse_ndof = 2*(c_nelx+1)*(c_nely+1)
	# Allocate design variables (as array), initialize and allocate sens.
	x=volfrac * np.ones(nely*nelx,dtype=float)
	xold=x.copy()
	xPhys=x.copy()
	g=0 # must be initialized to use the NGuyen/Paulino OC approach
	dc=np.zeros((nely,nelx), dtype=float)
	# FE: Build the index vectors for the for coo matrix format.
	KE=lk()
	edofMat=np.zeros((nelx*nely,8),dtype=int)
	for elx in range(nelx):
		for ely in range(nely):
			el = ely+elx*nely
			n1=(nely+1)*elx+ely
			n2=(nely+1)*(elx+1)+ely
			edofMat[el,:]=np.array([2*n1+2, 2*n1+3, 2*n2+2, 2*n2+3,2*n2, 2*n2+1, 2*n1, 2*n1+1])
	# Construct the index pointers for the coo format
	iK = np.kron(edofMat,np.ones((8,1))).flatten()
	jK = np.kron(edofMat,np.ones((1,8))).flatten()    
 
 
	coarse_edofMat=np.zeros((c_nelx*c_nely,8),dtype=int)
	for elx in range(c_nelx):
		for ely in range(c_nely):
			el = ely+elx*c_nely
			n1=(c_nely+1)*elx+ely
			n2=(c_nely+1)*(elx+1)+ely
			coarse_edofMat[el,:]=np.array([2*n1+2, 2*n1+3, 2*n2+2, 2*n2+3,2*n2, 2*n2+1, 2*n1, 2*n1+1])
	# Construct the index pointers for the coo format
	coarse_iK = np.kron(coarse_edofMat,np.ones((8,1))).flatten()
	coarse_jK = np.kron(coarse_edofMat,np.ones((1,8))).flatten()    
 
 
 
	# Filter: Build (and assemble) the index+data vectors for the coo matrix format
	nfilter=int(nelx*nely*((2*(np.ceil(rmin)-1)+1)**2))
	iH = np.zeros(nfilter)
	jH = np.zeros(nfilter)
	sH = np.zeros(nfilter)
	cc=0
	for i in range(nelx):
		for j in range(nely):
			row=i*nely+j
			kk1=int(np.maximum(i-(np.ceil(rmin)-1),0))
			kk2=int(np.minimum(i+np.ceil(rmin),nelx))
			ll1=int(np.maximum(j-(np.ceil(rmin)-1),0))
			ll2=int(np.minimum(j+np.ceil(rmin),nely))
			for k in range(kk1,kk2):
				for l in range(ll1,ll2):
					col=k*nely+l
					fac=rmin-np.sqrt(((i-k)*(i-k)+(j-l)*(j-l)))
					iH[cc]=row
					jH[cc]=col
					sH[cc]=np.maximum(0.0,fac)
					cc=cc+1
	# Finalize assembly and convert to csc format
	H=coo_matrix((sH,(iH,jH)),shape=(nelx*nely,nelx*nely)).tocsc()	
	Hs=H.sum(1)
	# BC's and support
	# dofs=np.arange(2*(nelx+1)*(nely+1))
	# fixed=np.union1d(dofs[0:2*(nely+1):2],np.array([2*(nelx+1)*(nely+1)-1]))
	# free=np.setdiff1d(dofs,fixed)
 
	coarse_dofs=np.arange(2*(c_nelx+1)*(c_nely+1))
	coarse_fixed=np.union1d(coarse_dofs[0:2*(c_nely+1):2],np.array([2*(c_nelx+1)*(c_nely+1)-1]))
	coarse_free=np.setdiff1d(coarse_dofs,coarse_fixed)
 
	# Solution and RHS vectors
	# f=np.zeros((ndof,1))
	# u=np.zeros((ndof,1))

	c_f=np.zeros((coarse_ndof,1))
	c_u=np.zeros((coarse_ndof,1))
	# Set load
	# f[1,0]=-1
 	
	c_f[1,0]=-1
	# Initialize plot and plot the initial design
	plt.ion() # Ensure that redrawing is possible
	fig,ax = plt.subplots()
	im = ax.imshow(-xPhys.reshape((nelx,nely)).T, cmap='gray',\
	interpolation='none',norm=colors.Normalize(vmin=-1,vmax=0))
	fig.show()
   	# Set loop counter and gradient vectors 
	loop=0
	change=1
	dv = np.ones(nely*nelx)
	dc = np.ones(nely*nelx)
	ce = np.ones(nely*nelx)
	while change>0.01 and loop<2000:
		loop=loop+1
		# Setup and solve FE problem
		coarse_xPhys=xPhys[12::25]
		sK=((KE.flatten()[np.newaxis]).T*(Emin+(coarse_xPhys)**penal*(Emax-Emin))).flatten(order='F')
		K = coo_matrix((sK,(coarse_iK,coarse_jK)),shape=(coarse_ndof,coarse_ndof)).tocsc()
		# Remove constrained dofs from matrix
		K = K[coarse_free,:][:,coarse_free]
		# Solve coarse situation 
		c_u[coarse_free,0]=spsolve(K,c_f[coarse_free,0])  
		# Predict fine situation
		u=pred_net(c_u,xPhys,c_nelx,c_nely,m,'checkpoints/ANN_mod1/ANN_mod1_opt.pt')
  
		# print(f.shape, f)
		# print(K.shape, K)
		# print(f[free,0])
		# print(u.shape, u)

		# Objective and sensitivity
		ce[:] = (np.dot(u[edofMat].reshape(nelx*nely,8),KE) * u[edofMat].reshape(nelx*nely,8) ).sum(1)
		obj=( (Emin+xPhys**penal*(Emax-Emin))*ce ).sum()
		dc[:]=(-penal*xPhys**(penal-1)*(Emax-Emin))*ce
		dv[:] = np.ones(nely*nelx)
		# Sensitivity filtering:
		if ft==0:
			dc[:] = np.asarray((H*(x*dc))[np.newaxis].T/Hs)[:,0] / np.maximum(0.001,x)
		elif ft==1:
			dc[:] = np.asarray(H*(dc[np.newaxis].T/Hs))[:,0]
			dv[:] = np.asarray(H*(dv[np.newaxis].T/Hs))[:,0]
		# Optimality criteria
		xold[:]=x
		(x[:],g)=oc(nelx,nely,x,volfrac,dc,dv,g)
		# Filter design variables
		if ft==0:   xPhys[:]=x
		elif ft==1:	xPhys[:]=np.asarray(H*x[np.newaxis].T/Hs)[:,0]
		# Compute the change by the inf. norm
		change=np.linalg.norm(x.reshape(nelx*nely,1)-xold.reshape(nelx*nely,1),np.inf)
		# Plot to screen
		im.set_array(-xPhys.reshape((nelx,nely)).T)
		fig.canvas.draw()
		# Write iteration history to screen (req. Python 2.6 or newer)
		print("it.: {0} , obj.: {1:.3f} Vol.: {2:.3f}, ch.: {3:.3f}".format(\
					loop,obj,(g+volfrac*nelx*nely)/(nelx*nely),change))
		
  
		np.save('results/top88_' + mod_idx + '_xPhys_' + str(nelx) + '_' + str(nely) + '.npy', xPhys.reshape((nelx,nely)).T)
		np.save('results/top88_' + mod_idx + '_u_' + str(nelx) + '_' + str(nely) + '.npy', u)
		plt.savefig('results/top88_' + mod_idx + '_img_' + str(nelx) + '_' + str(nely) + '.jpg')
  
		# plt.show()
  
	print(u.reshape(nelx+1,nely+1,2))

	# Make sure the plot stays and that the shell remains	
	np.save('results/top88_' + mod_idx + '_xPhys_' + str(nelx) + '_' + str(nely) + '.npy', xPhys.reshape((nelx,nely)).T)
	np.save('results/top88_' + mod_idx + '_u_' + str(nelx) + '_' + str(nely) + '.npy', u)
	plt.savefig('results/top88_' + mod_idx + '_img_' + str(nelx) + '_' + str(nely) + '.jpg')
	plt.show()
	print("Press any key...")
#element stiffness matrix
def lk():
	E=1
	nu=0.3
	k=np.array([1/2-nu/6,1/8+nu/8,-1/4-nu/12,-1/8+3*nu/8,-1/4+nu/12,-1/8-nu/8,nu/6,1/8-3*nu/8])
	KE = E/(1-nu**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]] ]);
	return (KE)
# Optimality criterion
def oc(nelx,nely,x,volfrac,dc,dv,g):
	l1=0
	l2=1e9
	move=0.2
	# reshape to perform vector operations
	xnew=np.zeros(nelx*nely)
	while (l2-l1)/(l1+l2)>1e-3:
		lmid=0.5*(l2+l1)
		xnew[:]= np.maximum(0.0,np.maximum(x-move,np.minimum(1.0,np.minimum(x+move,x*np.sqrt(-dc/dv/lmid)))))
		gt=g+np.sum((dv*(xnew-x)))
		if gt>0 :
			l1=lmid
		else:
			l2=lmid
	return (xnew,gt)


def pred_net(coarse_u,fine_x,coarse_nelx,coarse_nely,m, model_load_path, standard = False, device = 0):
	coarse_density=np.zeros(shape=(coarse_nely*coarse_nelx,m*m))
	fine_x=fine_x.reshape(coarse_nely*m,coarse_nelx*m)
	for ely in range(coarse_nely):
		for elx in range(coarse_nelx):
			coarse_density[elx + ely * m] = fine_x[ely * m : (ely + 1) * m, elx * m : (elx + 1) * m].flatten()
	print(coarse_density.shape)

 
	global_displace = coarse_u.reshape(coarse_nelx+1,coarse_nely+1,2)
	global_displace = np.dstack((global_displace[:,:,0].T, global_displace[:,:,1].T))
	coarse_displace=np.zeros(shape=(coarse_nely*coarse_nelx,4,2))
	for ely in range(coarse_nely):
		for elx in range(coarse_nelx):
			coarse_displace[elx + ely][0] = global_displace[ely, elx, :]
			coarse_displace[elx + ely][1] = global_displace[ely, (elx+1), :]
			coarse_displace[elx + ely][2] = global_displace[(ely+1), elx, :]
			coarse_displace[elx + ely][3] = global_displace[(ely+1), (elx+1), :]
	print(coarse_displace.shape)
 
	X = np.hstack((coarse_density[:,:] , coarse_displace[:,:,0] , coarse_displace[:,:,1]))
	model = torch.load(model_load_path)
	if standard:
		X = standardization(X)
	device = f'cuda:{device}' if torch.cuda.is_available() else 'cpu'
	X = torch.from_numpy(X).type(torch.float32).to(device)
	pred=model(X).cpu().detach().numpy()
	
	# print(pred)
	u_reconstructed = np.zeros(shape=(coarse_nely*m+1, coarse_nelx*m+1, 2))
	for ely in range(coarse_nely):
		for elx in range(coarse_nelx):
			u_reconstructed[ely*m : (ely+1)*m+1, elx*m : (elx+1)*m+1, :] = pred[elx + ely * m].reshape((m+1, m+1, 2))
	# u_reconstructed = u_reconstructed.reshape(coarse_nely*m+1, coarse_nelx*m+1,2)
	u_reconstructed = np.dstack((u_reconstructed[:,:,0].T, u_reconstructed[:,:,1].T))
	u_reconstructed=u_reconstructed.flatten()
	
	return u_reconstructed


 

# The real main driver    
if __name__ == "__main__":
	mod_idx='test1'
	m=5
	nelx=180
	nely=60
	volfrac=0.4
	rmin=5.4
	penal=3.0
	ft=1 # ft==0 -> sens, ft==1 -> dens
	top_EMsFEA(nelx,nely,volfrac,penal,rmin,ft,mod_idx,m)