update with new data processor

README.md

@@ -33,6 +33,7 @@ python test.py
 
 ```
 python topopt_EMsFEA.py
+# 参数详见options/topopt_options.py
 ```
 
 ## 数据集
@@ -40,8 +41,6 @@ python topopt_EMsFEA.py
 >
 > Download from:
 > 
-> http://118.195.195.192:3000/GyeongYun/EMsFEA-net/raw/branch/resources/checkpoints.zip
->
 > http://118.195.195.192:3000/GyeongYun/EMsFEA-net/raw/branch/resources/datasets.zip
 
 mod1: ![](doc/mod1.jpg)

models/ANN.py

@@ -3,42 +3,36 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 class ANN_Model(nn.Module):
-    def __init__(self,input_features=8,out_features=72):
+    def __init__(self,in_dim=25,l1=36,l2=36*2,l3=36*2,l4=36*4,l5=36*4,l6=36*8,l7=36*8,l8=36*16,out_dim=36*2*4*2):
         super().__init__()        
-        self.fc1=nn.Linear(input_features,12)
-        self.fc2=nn.Linear(12,16)
-        self.fc3=nn.Linear(16,20)
-        self.fc4=nn.Linear(20,25)
+        self.fc1=nn.Linear(in_dim,l1)
+        self.fc2=nn.Linear(l1,l2)
+        self.fc3=nn.Linear(l2,l3)
+        self.fc4=nn.Linear(l3,l4)
         
-        self.fc5=nn.Linear(50,60)
-        self.fc6=nn.Linear(60,70)
-        self.fc7=nn.Linear(70,80)
-        self.fc8=nn.Linear(80,90)
-        self.fc9=nn.Linear(90,100)
-        self.fc10=nn.Linear(100,90)
-        self.fc11=nn.Linear(90,80)
+        self.fc5=nn.Linear(l4,l5)
+        self.fc6=nn.Linear(l5,l6)
+        self.fc7=nn.Linear(l6,l7)
+        self.fc8=nn.Linear(l7,l8)
 
         
-        self.out=nn.Linear(80,out_features)   
+        self.out=nn.Linear(l8,out_dim) # -> 576
 
 
     def forward(self,x):
-        density=x[:,:25].reshape(x.shape[0],25)
-        displace = x[:,25:]
-        x = F.relu(self.fc1(displace))
+        density = x[:25]
+        displace = x[25:]
+        
+        x = F.relu(self.fc1(density))
         x = F.relu(self.fc2(x))
         x = F.relu(self.fc3(x))
         x = F.relu(self.fc4(x))
-        x = torch.hstack((density,x))
+        # x = torch.hstack((density,x))
 
         x = F.relu(self.fc5(x))
         x = F.relu(self.fc6(x))
         x = F.relu(self.fc7(x))
         x = F.relu(self.fc8(x))
-        x = F.relu(self.fc9(x))
-        x = F.relu(self.fc10(x))
-        x = F.relu(self.fc11(x))
-
 
         x = self.out(x)
         

options/base_options.py

@@ -1,5 +1,6 @@
 import argparse
 import os
+import time
 from utils import utils
 import torch
 
@@ -24,6 +25,7 @@ class BaseOptions():
         parser.add_argument('--device', type=str, default='cuda:0', help='generate device with gpu_id and usable user device')
         parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', help='models are saved here')
         parser.add_argument('--ms_ratio', type=int, default=5, help='multiscale ratio')
+        parser.add_argument('--results_dir', type=str, default='./results', help='saves results here.')
 
         # model parameters
         parser.add_argument('--model', type=str, default='ANN', help='chooses which model to use. [ANN | CNN | AutoEncoder]')
@@ -31,9 +33,9 @@ class BaseOptions():
         # parser.add_argument('--resolution', type=str, default='180_60', help='data resolution. nelx_nely here')
         parser.add_argument('--nelx', type=int, default=180, help='num of elements on x-axis')
         parser.add_argument('--nely', type=int, default=60, help='num of elements on y-axis')
-        parser.add_argument('--nelz', type=int, default=1, help='num of elements on z-axis')
+        parser.add_argument('--nelz', type=int, default=0, help='num of elements on z-axis')
         parser.add_argument('--dimension', type=int, default=2, help='dimension of dataset models')
-        parser.add_argument('--is_standard', type=bool, default=False, help='whether need standardization or not')
+        parser.add_argument('--is_standard', type=bool, default=True, help='whether need standardization or not')
 
         # additional parameters
         parser.add_argument('--epoch', type=str, default='latest', help='which epoch to load? set to latest to use latest cached model')
@@ -80,12 +82,16 @@ class BaseOptions():
         print(message)
 
         # save to the disk
-        expr_dir = os.path.join(opt.checkpoints_dir, opt.model+'_'+opt.mod)
-        utils.mkdir(expr_dir)
-        file_name = os.path.join(expr_dir, '{}_opt.txt'.format(opt.phase))
-        with open(file_name, 'wt') as opt_file:
-            opt_file.write(message)
-            opt_file.write('\n')
+        if opt.isTrain:
+            curr=time.strftime('%y%m%d-%H%M%S')
+            expr_dir = os.path.join(opt.checkpoints_dir, opt.model+'_'+opt.mod+'_'+str(curr))
+            opt.expr_dir = expr_dir
+            
+            utils.mkdir(expr_dir)
+            file_name = os.path.join(expr_dir, '{}_opt.txt'.format(opt.phase))
+            with open(file_name, 'wt') as opt_file:
+                opt_file.write(message)
+                opt_file.write('\n')
 
     def parse(self):
         """Parse our options, create checkpoints directory suffix, and set up gpu device."""

options/test_options.py

@@ -9,10 +9,9 @@ class TestOptions(BaseOptions):
 
     def initialize(self, parser):
         parser = BaseOptions.initialize(self, parser)  # define shared options
-        parser.add_argument('--results_dir', type=str, default='./results', help='saves results here.')
         parser.add_argument('--phase', type=str, default='test', help='train, val, test, etc')
-        parser.add_argument('--mod', type=str, default='mod3', help='chooses which dataset model for test. mod1....')
-        parser.add_argument('--pretrained_model_path', type=str, default='./checkpoints/ANN_mod1/ANN_mod1_opt.pt', help='pretrained model file load path')
+        parser.add_argument('--mod', type=str, default='mod2', help='chooses which dataset model for test. mod1....')
+        parser.add_argument('--pretrained_model_path', type=str, default='./checkpoints/ANN_mod1_231224-222338/ANN_mod1_opt.pt', help='pretrained model file load path')
 
 
         self.isTrain = False

options/topopt_options.py

@@ -0,0 +1,26 @@
+from .base_options import BaseOptions
+
+
+class TopoptOption(BaseOptions):
+    """This class includes test options.
+
+    It also includes shared options defined in BaseOptions.
+    """
+
+    def initialize(self, parser):
+        parser = BaseOptions.initialize(self, parser)  # define shared options
+        parser.add_argument('--phase', type=str, default='topopt', help='train, val, test, etc')
+        parser.add_argument('--pretrained_model_path', type=str, default='./checkpoints/ANN_mod1_231224-222338/ANN_mod1_opt.pt', help='pretrained model file load path')
+        parser.add_argument('--mod_idx', type=str, default='mod1', help='mod_idx for identify save path') 
+
+        parser.add_argument('--nelx_to', type=int, default=180, help='num of elements on x-axis')
+        parser.add_argument('--nely_to', type=int, default=60, help='num of elements on y-axis')
+        parser.add_argument('--ms_ratio_to', type=int, default=5, help='multiscale ratio')
+        parser.add_argument('--volfrac', type=float, default=0.4, help='volfrac')
+        parser.add_argument('--rmin', type=float, default=5.4, help='rmin')
+        parser.add_argument('--penal', type=float, default=3.0, help='penal')
+        parser.add_argument('--ft', type=int, default=1, help='ft')
+
+
+        self.isTrain = False
+        return parser

options/train_options.py

@@ -15,9 +15,9 @@ class TrainOptions(BaseOptions):
         parser.add_argument('--phase', type=str, default='train', help='train, val, test, etc')
 
         # training parameters
-        parser.add_argument('--epochs', type=int, default=10000, help='number of epochs')
-        parser.add_argument('--lr', type=float, default=0.001, help='initial learning rate for adam')
-        parser.add_argument('--mod', type=str, default='mod2', help='chooses which dataset model for train. mod1....')
+        parser.add_argument('--epochs', type=int, default=200, help='number of epochs')
+        parser.add_argument('--lr', type=float, default=1e-5, help='initial learning rate for adam')
+        parser.add_argument('--mod', type=str, default='mod1', help='chooses which dataset model for train. mod1....')
 
         self.isTrain = True
         return parser

test.py

@@ -7,19 +7,22 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 from utils.data_standardizer import standardization
-from utils.data_loader import data_loader
+from utils.data_loader import data_loader_new
 from options.test_options import TestOptions
-
+from utils.visualization import surf_plot
 
 def test(X, opt):
-    # OLD: model_load_path, X, standard = False, device = 0
-    if opt.is_standard:
-        X = standardization(X)
-        
     X_test=torch.from_numpy(X).type(torch.float32).to(opt.device)
 
     model = torch.load(opt.pretrained_model_path)
-    return model(X_test)
+    
+    N=(opt.ms_ratio+1)**2 * 2
+    pred=torch.zeros(X_test.shape[0], N)
+    for batch_idx, data_batch in enumerate(X_test):
+        pred_ShapeFunction=model(data_batch)
+        pred[batch_idx,:]=pred_ShapeFunction.reshape(N,8) @ data_batch[25:]
+    
+    return pred
 
 
 if __name__=='__main__':
@@ -27,18 +30,27 @@ if __name__=='__main__':
     opt = TestOptions().parse()
     
     # Load datasets, mod2 as default
-    global_density, global_displace, coarse_density, coarse_displace, fine_displace = data_loader(opt)
-    X = np.hstack((coarse_density[:,:] , coarse_displace[:,:,0] , coarse_displace[:,:,1]))
-    Y = fine_displace[:,:]
-    
+    m=opt.ms_ratio
+    c_nelx=int(opt.nelx/m)
+    c_nely=int(opt.nely/m)
+    c_N=c_nelx*c_nely
+    global_density, global_displace, coarse_density, coarse_displace, fine_displace = data_loader_new(opt)
+    # X = np.hstack((coarse_density.reshape(c_N,m*m), coarse_displace.reshape(c_N,2,2,2)[:,:,:,0].reshape(c_N,4), coarse_displace.reshape(c_N,2,2,2)[:,:,:,1].reshape(c_N,4)))
+    X = np.hstack((coarse_density.reshape(c_N,m*m), coarse_displace.reshape(c_N,8)))
+    Y = fine_displace.reshape(c_N,(m+1)**2*2)
+    if opt.is_standard:
+        X = standardization(X)
+        Y = standardization(Y)
+
     # Predict 
     pred = test(X, opt)
-
+    pred.to('cpu')
+        
     # Set loss function
     loss_function = nn.MSELoss()
     # Calculate loss
     pred_loss=[]
-    Y_test = torch.from_numpy(Y).type(torch.float32).to(opt.device)
+    Y_test = torch.from_numpy(Y).type(torch.float32).to('cpu')
     for i in range(pred.shape[0]):
         pred_loss.append(loss_function(pred[i,:],Y_test[i,:]).item())
     
@@ -58,3 +70,8 @@ if __name__=='__main__':
     plt.title("Show loss value in grid") 
     plt.savefig(os.path.join(opt.results_dir, 'test_loss_in_grid.png'))
     plt.show()
+    
+    # Plot every mesh displacement for comparation
+    # for i in range(pred.shape[0]):
+    #     surf_plot(pred[i].detach().numpy(),6,6,os.path.join(opt.results_dir, 'meshes', 'test_pred_mesh'+str(i)+'.png'),'u')
+    #     surf_plot(Y_test[i].detach().numpy(),6,6,os.path.join(opt.results_dir, 'meshes','test_GT_mesh'+str(i)+'.png'),'u')

topopt_EMsFEA.py

@@ -1,5 +1,6 @@
-from __future__ import division
+# from __future__ import division
 import numpy as np
+import os
 from scipy.sparse import coo_matrix
 from scipy.sparse.linalg import spsolve
 from matplotlib import colors
@@ -11,251 +12,274 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 from utils.data_standardizer import standardization
-from utils.data_loader import data_loader
+from utils.mesh_reshape import Ms_u_reshape
+from options.topopt_options import TopoptOption
 
-
-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()    
+def top_EMsFEA(opt):
+    mod_idx=opt.mod_idx
+    m=opt.ms_ratio_to
+    nelx=opt.nelx_to
+    nely=opt.nely_to
+    volfrac=opt.volfrac
+    rmin=opt.rmin
+    penal=opt.penal
+    ft=opt.ft # ft==0 -> sens, ft==1 -> dens
+    
+    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()    
+    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)
+    # 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)
+    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))
+    # 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')
+    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,opt)
+        
+        # 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)
+        # 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))
-		
+        # 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')
+        np.save('results/EMsNetTop_' + mod_idx + '_xPhys_' + str(nelx) + '_' + str(nely) + '.npy', xPhys.reshape((nelx,nely)).T)
+        np.save('results/EMsNetTop_' + mod_idx + '_u_' + str(nelx) + '_' + str(nely) + '.npy', u)
+        plt.savefig('results/EMsNetTop_' + mod_idx + '_img_' + str(nelx) + '_' + str(nely) + '.jpg')
   
-		# plt.show()
+        # plt.show()
   
-	print(u.reshape(nelx+1,nely+1,2))
+    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...")
+    # Make sure the plot stays and that the shell remains	
+    np.save('results/EMsNetTop_' + mod_idx + '_xPhys_' + str(nelx) + '_' + str(nely) + '.npy', xPhys.reshape((nelx,nely)).T)
+    np.save('results/EMsNetTop_' + mod_idx + '_u_' + str(nelx) + '_' + str(nely) + '.npy', u)
+    plt.savefig('results/EMsNetTop_' + 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)
+    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)
+    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)
+def pred_net(coarse_u,global_x,opt):
+    m=opt.ms_ratio_to
+    nelx=opt.nelx_to
+    nely=opt.nely_to
+    coarse_nelx=int(nelx/m)
+    coarse_nely=int(nely/m)
+    c_N=coarse_nelx*coarse_nely
+    N=(opt.ms_ratio_to+1)**2 * 2
+    
+    # Generate coarse mesh density
+    global_density=global_x.reshape(nelx,nely) # -> (nelx , nely)
+    coarse_density = np.lib.stride_tricks.as_strided(
+        global_density,
+        shape=(coarse_nelx, coarse_nely, m, m),  
+        strides=global_density.itemsize * np.array([nely*m, m, nely, 1])
+    )
 
- 
-	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
+    # Generate coarse mesh displacement
+    coarse_displace= np.lib.stride_tricks.as_strided(
+        coarse_u,
+        shape=(coarse_nelx, coarse_nely, 2, 2, 2), 
+        strides=coarse_u.itemsize * np.array([(coarse_nely+1)*2, 1*2, (coarse_nely+1)*2, 1*2, 1])
+    )
+    
+    # data preprocess
+    X = np.hstack((coarse_density.reshape(c_N,m*m), coarse_displace.reshape(c_N,8)))
+    if opt.is_standard:
+        X = standardization(X)
+    X=torch.from_numpy(X).type(torch.float32).to(opt.device)
 
+    # predict
+    model = torch.load(opt.pretrained_model_path)
+    pred=torch.zeros(X.shape[0], N)
+    for batch_idx, data_batch in enumerate(X):
+        pred_ShapeFunction=model(data_batch)
+        pred[batch_idx,:]=pred_ShapeFunction.reshape(N,8) @ data_batch[25:]
+    pred=pred.to('cpu').detach().numpy()
+    pred=Ms_u_reshape(pred, coarse_nelx, coarse_nely, m)
+    pred=pred.reshape((nelx+1)*(nely+1)*2,1)
+    
+    return pred
 
  
 
 # 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)
+    # Load parmetaers
+    opt = TopoptOption().parse()
+    # 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(opt)
+    
+    # u=np.load('./results/coarse_u.npy')
+    # x=np.load('./results/global_x.npy')
+    # pred_net(u,x,opt)
+

train.py

@@ -8,7 +8,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 from utils.data_standardizer import standardization
-from utils.data_loader import data_loader
+from utils.data_loader import data_loader_new
 from options.train_options import TrainOptions
 
 from models.ANN import ANN_Model
@@ -33,25 +33,34 @@ def train(X, Y, opt):
     loss_function = nn.MSELoss()
     
     # Set adam optimizer
-    optimizer=torch.optim.Adam(model.parameters(),lr=opt.lr) # ANN 学习率最好0.001 左右（无归一化）
+    optimizer=torch.optim.Adam(model.parameters(),lr=opt.lr)
     
     # Train
     start_time=time.time()
     losses=[]
-    for i in range(opt.epochs):
-        pred = model.forward(X_train)
-        loss=loss_function(pred,Y_train)
-        # loss.requires_grad_(True)
+    loss=0
+    N=(opt.ms_ratio+1)**2 * 2
+    for epoch in range(opt.epochs):
+        for batch_idx, data_batch in enumerate(X_train):
+            model.train() # 启用 batch normalization 和 dropout
+
+            pred_ShapeFunction = model(data_batch)
+            pred_U=pred_ShapeFunction.reshape(N,8) @ data_batch[25:]
+            loss=loss_function(pred_U,Y_train[batch_idx,:])
+            # print(loss.item())
+            # loss.requires_grad_(True)
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
         losses.append(loss.cpu().detach().numpy())
-        if i%(opt.epochs/10)==1:
-            print("Epoch number: {} and the loss : {}".format(i,loss.item()))
-        optimizer.zero_grad()
-        loss.backward()
-        optimizer.step()
+        if epoch%(opt.epochs/20)==1:
+            print("Epoch number: {} and the loss : {}".format(epoch,loss.item()))
     print(time.time()-start_time)
     
-    # save trained model, mkdir opt has done in options/base_options.py
-    save_path=os.path.join(opt.checkpoints_dir, opt.model+'_'+opt.mod, opt.model+'_'+opt.mod+'_opt.pt')
+    # save trained model, mkdir opreate has done in options/base_options.py
+    save_path=os.path.join(opt.expr_dir, opt.model+'_'+opt.mod+'_opt.pt')
     torch.save(model, save_path)
     
     return losses
@@ -60,11 +69,18 @@ def train(X, Y, opt):
 if __name__=='__main__':
     # Load parmetaers
     opt = TrainOptions().parse()
-    
+    save_path=os.path.join(opt.expr_dir, opt.model+'_'+opt.mod+'_opt.pt')
+
     # Load datasets, mod1 as default
-    global_density, global_displace, coarse_density, coarse_displace, fine_displace = data_loader(opt)
-    X = np.hstack((coarse_density[:,:] , coarse_displace[:,:,0] , coarse_displace[:,:,1]))
-    Y = fine_displace[:,:]
+    m=opt.ms_ratio
+    c_nelx=int(opt.nelx/m)
+    c_nely=int(opt.nely/m)
+    c_N=c_nelx*c_nely
+    
+    global_density, global_displace, coarse_density, coarse_displace, fine_displace = data_loader_new(opt)
+    # X = np.hstack((coarse_density.reshape(c_N,m*m), coarse_displace.reshape(c_N,2,2,2)[:,:,:,0].reshape(c_N,4), coarse_displace.reshape(c_N,2,2,2)[:,:,:,1].reshape(c_N,4)))
+    X = np.hstack((coarse_density.reshape(c_N,m*m), coarse_displace.reshape(c_N,8)))
+    Y = fine_displace.reshape(c_N,(m+1)**2*2)
 
     # Train
     losses = train(X, Y, opt)
@@ -73,4 +89,5 @@ if __name__=='__main__':
     plt.plot(range(opt.epochs),losses)
     plt.ylabel('Loss')
     plt.xlabel('Epoch')
+    plt.savefig(os.path.join(opt.results_dir, 'train_losses.png'))
     plt.show()

utils/data_loader.py

@@ -1,6 +1,54 @@
 import numpy as np
 import os
 
+def data_loader_new(opt):
+    nelx=opt.nelx
+    nely=opt.nely
+    m=opt.ms_ratio
+    coarse_nelx = int(nelx/m)
+    coarse_nely = int(nely/m)
+    
+    # './datasets/train/180_60/u_OR_xPhys/mod1.npy' as default
+    density_load_path = os.path.join(opt.dataroot, opt.phase, str(opt.nelx)+'_'+str(opt.nely), 'xPhys', opt.mod+'.npy')
+    displace_load_path = os.path.join(opt.dataroot, opt.phase, str(opt.nelx)+'_'+str(opt.nely), 'u', opt.mod+'.npy')
+
+    global_density = np.load(density_load_path) # (nely , nelx)
+    global_displace = np.load(displace_load_path) # ( (nely+1)*(nelx+1)*2 , 1 )
+
+    # 有意义的情况：
+    global_density=global_density.reshape(nely,nelx).T # -> (nelx , nely)
+    global_displace=global_displace.reshape(nelx+1,nely+1,2) # -> (nelx+1), (nely+1), 2 
+    print(global_density.shape)
+    print(global_displace.shape)
+
+    
+    # Generate coarse mesh density
+    coarse_density = np.lib.stride_tricks.as_strided(
+        global_density,
+        shape=(coarse_nelx, coarse_nely, m, m),  # 要输出矩阵的 shape
+        strides=global_density.itemsize * np.array([nely*m, m, nely, 1])
+    )
+    print(coarse_density.shape)
+    
+    # Generate coarse mesh displacement
+    coarse_displace= np.lib.stride_tricks.as_strided(
+        global_displace,
+        shape=(coarse_nelx, coarse_nely, 2, 2, 2),  # 要输出矩阵的 shape
+        strides=global_displace.itemsize * np.array([(nely+1)*m*2, m*2, (nely+1)*m*2, m*2, 1])
+    )
+    print(coarse_displace.shape)
+    
+    # Generate fine mesh displacement
+    fine_displace= np.lib.stride_tricks.as_strided(
+        global_displace,
+        shape=(coarse_nelx, coarse_nely, m+1, m+1, 2),  # 要输出矩阵的 shape
+        strides=global_displace.itemsize * np.array([(nely+1)*m*2, m*2, (nely+1)*2, 1*2, 1])
+    )
+    print(fine_displace.shape)
+    
+    return global_density, global_displace, coarse_density, coarse_displace, fine_displace
+
+
 def data_loader(opt):
     # Load datasets
     
@@ -49,6 +97,7 @@ def data_loader(opt):
     return global_density, global_displace, coarse_density, coarse_displace, fine_displace
     
 if __name__=='__main__':
-    from options.train_options import TrainOptions
-    opt = TrainOptions().parse()
-    data_loader(opt)
+    global_density, global_displace, coarse_density, coarse_displace, fine_displace=data_loader_new('datasets/train/180_60/xPhys/mod1.npy', 'datasets/train/180_60/u/mod1.npy')
+    print(global_displace[:10,:10,0])
+    print(fine_displace[33,11,:,:,0])
+    print(coarse_displace[33,11,:,:,0])

utils/mesh_reshape.py

@@ -0,0 +1,28 @@
+import numpy as np
+
+def Ms_u_reshape(u_data, coarse_nelx, coarse_nely, m):
+    nelx=coarse_nelx*m
+    nely=coarse_nely*m
+    
+    u_data = u_data.reshape(coarse_nelx,coarse_nely,m+1,m+1,2)
+    u_data = u_data.swapaxes(1,2).reshape(coarse_nelx*(m+1), coarse_nely*(m+1), 2)  
+    
+    idx_x=np.arange(coarse_nelx*(m+1))[::m+1]
+    idx_x=np.delete(idx_x,0)
+    idx_x=np.delete(np.arange(coarse_nelx*(m+1)),idx_x)
+    idx_y=np.arange(coarse_nely*(m+1))[::m+1]
+    idx_y=np.delete(idx_y,0)
+    idx_y=np.delete(np.arange(coarse_nely*(m+1)),idx_y)
+
+    return u_data[idx_x.reshape(nelx+1,1),idx_y.reshape(1,nely+1)]
+
+
+if __name__=='__main__':
+    pred=np.load('results/pred.npy')
+    u=np.load('datasets/train/180_60/u/mod2.npy')
+    
+    print(u.shape)
+    print(pred.shape)
+
+    recv_u = Ms_u_reshape(pred, 36, 12, 5)
+    print(recv_u-u.reshape(181,61,2))

utils/topopt_88.py

@@ -121,16 +121,16 @@ def top88(nelx,nely,volfrac,penal,rmin,ft,mod_idx):
 					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')
+		# 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 + '_xPhys_' + str(nelx) + '_' + str(nely) + '.npy', xPhys)
 	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()
@@ -168,9 +168,9 @@ def oc(nelx,nely,x,volfrac,dc,dv,g):
 # The real main driver    
 if __name__ == "__main__":
 	# Default input parameters
-	mod_idx='mod4'
-	nelx=30
-	nely=10
+	mod_idx='py88'
+	nelx=180
+	nely=60
 	volfrac=0.4
 	rmin=5.4
 	penal=3.0

utils/visualization.py

@@ -0,0 +1,35 @@
+from mpl_toolkits.mplot3d import Axes3D
+import matplotlib.pyplot as plt
+import numpy as np
+
+def surf_plot(dp,ny,nx,save_path,UorV='sqrt',non_dp=False):
+    fig = plt.figure()
+    ax = fig.add_subplot(projection='3d')
+
+    # Make data
+    # 生成网格数据
+    X, Y = np.meshgrid(np.arange(nx), np.arange(ny))
+
+    if non_dp:
+        Z = Z_x = Z_y = dp.reshape(ny,nx)
+    else:
+        Z = dp.reshape(ny,nx,2)
+        Z_x = Z[:,:,0].reshape(ny,nx)
+        Z_y = Z[:,:,1].reshape(ny,nx)
+
+    if UorV == 'u':
+        ax.plot_surface(X, Y, Z_x, rstride = 1, cstride = 1, cmap='rainbow')
+    elif UorV == 'v':
+        ax.plot_surface(X, Y, Z_y, rstride = 1, cstride = 1, cmap='rainbow')
+    else:
+        ax.plot_surface(X, Y, Z_x*Z_x+Z_y*Z_y, rstride = 1, cstride = 1, cmap='rainbow')
+        
+    plt.savefig(save_path)
+    plt.show()
+
+
+if __name__=='__main__':
+    print('nothing')
+    # surf_plot(global_displace,nely+1,nelx+1,'u')
+    # surf_plot(global_displace,nely+1,nelx+1,'v')
+    # surf_plot(global_displace,nely+1,nelx+1,'sqrt')