├── README.md ├── static ├── asgard.jpeg ├── building.gif ├── causeway.gif ├── causeway.png ├── causeway_long.gif ├── eves.gif ├── points.gif ├── staircase.gif ├── stopt_design_domain.png ├── stopt_mbb_setup.png └── terminal.png ├── stopt_200.py └── stopt_240.py /README.md: -------------------------------------------------------------------------------- 1 | # Fast Structural Optimization in 200 Lines of Python 2 | 3 | [Blog post](https://greydanus.github.io/2022/05/08/structural-optimization/) | [Paper](https://arxiv.org/pdf/2205.08966.pdf) | [Colab notebook](https://bit.ly/394DUcL) 4 | 5 | ![causeway_long.gif](./static/causeway_long.gif) 6 | 7 | Structural optimization lets us design trusses, bridges, and buildings starting from the physics of elastic materials. Let's code it up, from scratch, in 200 lines (forty of these lines are whitespace and comments). 8 | 9 | There are two self-contained Python files in this directly. The first, `stopt_200.py`, provides the basic code you'll need in order to do structural optimization. The second, `stopt_240.py` adds LRU caching for NumPy arrays, which speeds up the code by 2x. Since caching isn't strictly necessary in a teaching example, we decided to separate it out. If you want to run this code on a larger problem, though, definitely use it! 10 | 11 | To run this code: clone the repo, `cd` into it, and run `python stopt_200.py ` or `python stopt_240.py ` 12 | 13 | ![terminal.png](./static/terminal.png) 14 | 15 | ## Scaling to larger designs 16 | 17 | Designing the eves of a gazebo: 18 | 19 | ![eves.gif](./static/eves.gif) 20 | 21 | Designing a causeway bridge: 22 | 23 | ![causeway.gif](./static/causeway.gif) 24 | 25 | Designing a structure that supports a grid of offset points: 26 | 27 | ![points.gif](./static/points.gif) 28 | 29 | Designing a staircase: 30 | 31 | ![staircase.gif](./static/staircase.gif) 32 | 33 | Designing a multistory building: 34 | 35 | ![building.gif](./static/building.gif) 36 | 37 | ## Dependencies 38 | 39 | This project uses `numpy` for dense matrix operations, `scipy` for sparse matrix operations, `matplotlib` for plotting, `autograd` for computing gradients, and `nlopt` for optimization. The only library that needs to be manually installed on Google Colab is `nlopt`. 40 | 41 | ## Final comments 42 | 43 | In sci-fi representations of the healthy cities of the future, we often find manmade structures that are well integrated with their natural surroundings. Sometimes we even see a convergence where nature has adapted to the city and the city has adapted to nature. The more decadent cities, on the other hand, tend to define themselves in opposition to the patterns of nature. Their architecture is more blocky and inorganic. Perhaps tools like structural optimization can help us build the healthy cities of the future – and steer clear of the decadent ones. 44 | 45 | ![asgard.jpeg](./static/asgard.jpeg) 46 | -------------------------------------------------------------------------------- /static/asgard.jpeg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/greydanus/structural_optimization/ce36c79979a956bb9853e7eb77510b9a7e922c85/static/asgard.jpeg -------------------------------------------------------------------------------- /static/building.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/greydanus/structural_optimization/ce36c79979a956bb9853e7eb77510b9a7e922c85/static/building.gif -------------------------------------------------------------------------------- /static/causeway.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/greydanus/structural_optimization/ce36c79979a956bb9853e7eb77510b9a7e922c85/static/causeway.gif -------------------------------------------------------------------------------- /static/causeway.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/greydanus/structural_optimization/ce36c79979a956bb9853e7eb77510b9a7e922c85/static/causeway.png -------------------------------------------------------------------------------- /static/causeway_long.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/greydanus/structural_optimization/ce36c79979a956bb9853e7eb77510b9a7e922c85/static/causeway_long.gif -------------------------------------------------------------------------------- /static/eves.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/greydanus/structural_optimization/ce36c79979a956bb9853e7eb77510b9a7e922c85/static/eves.gif -------------------------------------------------------------------------------- /static/points.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/greydanus/structural_optimization/ce36c79979a956bb9853e7eb77510b9a7e922c85/static/points.gif -------------------------------------------------------------------------------- /static/staircase.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/greydanus/structural_optimization/ce36c79979a956bb9853e7eb77510b9a7e922c85/static/staircase.gif -------------------------------------------------------------------------------- /static/stopt_design_domain.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/greydanus/structural_optimization/ce36c79979a956bb9853e7eb77510b9a7e922c85/static/stopt_design_domain.png -------------------------------------------------------------------------------- /static/stopt_mbb_setup.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/greydanus/structural_optimization/ce36c79979a956bb9853e7eb77510b9a7e922c85/static/stopt_mbb_setup.png -------------------------------------------------------------------------------- /static/terminal.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/greydanus/structural_optimization/ce36c79979a956bb9853e7eb77510b9a7e922c85/static/terminal.png -------------------------------------------------------------------------------- /stopt_200.py: -------------------------------------------------------------------------------- 1 | # A Tutorial on Structural Optimization | Sam Greydanus | 2022 2 | import time, nlopt # for optimization 3 | import numpy as np # for dense matrix ops 4 | import matplotlib.pyplot as plt # for plotting 5 | import autograd, autograd.core, autograd.extend, autograd.tracer # for adjoints 6 | import autograd.numpy as anp 7 | import scipy, scipy.ndimage, scipy.sparse, scipy.sparse.linalg # sparse matrices 8 | 9 | ##### Problem setup ##### 10 | class ObjectView(object): 11 | def __init__(self, d): self.__dict__ = d 12 | 13 | def get_args(normals, forces, density=0.4): # Manage the problem setup parameters 14 | width = normals.shape[0] - 1 15 | height = normals.shape[1] - 1 16 | fixdofs = np.flatnonzero(normals.ravel()) 17 | alldofs = np.arange(2 * (width + 1) * (height + 1)) 18 | freedofs = np.sort(list(set(alldofs) - set(fixdofs))) 19 | params = { 20 | # material properties 21 | 'young': 1, 'young_min': 1e-9, 'poisson': 0.3, 'g': 0, 22 | # constraints 23 | 'density': density, 'xmin': 0.001, 'xmax': 1.0, 24 | # input parameters 25 | 'nelx': width, 'nely': height, 'mask': 1, 'penal': 3.0, 'filter_width': 1, 26 | 'freedofs': freedofs, 'fixdofs': fixdofs, 'forces': forces.ravel(), 27 | # optimization parameters 28 | 'opt_steps': 80, 'print_every': 10} 29 | return ObjectView(params) 30 | 31 | def mbb_beam(width=80, height=25, density=0.4, y=1, x=0): # textbook beam example 32 | normals = np.zeros((width + 1, height + 1, 2)) 33 | normals[-1, -1, y] = 1 34 | normals[0, :, x] = 1 35 | forces = np.zeros((width + 1, height + 1, 2)) 36 | forces[0, 0, y] = -1 37 | return normals, forces, density 38 | 39 | def young_modulus(x, e_0, e_min, p=3): 40 | return e_min + x ** p * (e_0 - e_min) 41 | 42 | def physical_density(x, args, volume_contraint=False, use_filter=True): 43 | x = args.mask * x.reshape(args.nely, args.nelx) # reshape from 1D to 2D 44 | return gaussian_filter(x, args.filter_width) if use_filter else x # maybe filter 45 | 46 | def mean_density(x, args, volume_contraint=False, use_filter=True): 47 | return anp.mean(physical_density(x, args, volume_contraint, use_filter)) / anp.mean(args.mask) 48 | 49 | ##### Optimization objective + physics of elastic materials ##### 50 | def objective(x, args, volume_contraint=False, use_filter=True): 51 | kwargs = dict(penal=args.penal, e_min=args.young_min, e_0=args.young) 52 | x_phys = physical_density(x, args, volume_contraint=volume_contraint, use_filter=use_filter) 53 | ke = get_stiffness_matrix(args.young, args.poisson) # stiffness matrix 54 | u = displace(x_phys, ke, args.forces, args.freedofs, args.fixdofs, **kwargs) 55 | c = compliance(x_phys, u, ke, **kwargs) 56 | return c 57 | 58 | def compliance(x_phys, u, ke, *, penal=3, e_min=1e-9, e_0=1): 59 | nely, nelx = x_phys.shape 60 | ely, elx = anp.meshgrid(range(nely), range(nelx)) # x, y coords for the index map 61 | 62 | n1 = (nely+1)*(elx+0) + (ely+0) # nodes 63 | n2 = (nely+1)*(elx+1) + (ely+0) 64 | n3 = (nely+1)*(elx+1) + (ely+1) 65 | n4 = (nely+1)*(elx+0) + (ely+1) 66 | all_ixs = anp.array([2*n1, 2*n1+1, 2*n2, 2*n2+1, 2*n3, 2*n3+1, 2*n4, 2*n4+1]) 67 | u_selected = u[all_ixs] # select from u matrix 68 | 69 | ke_u = anp.einsum('ij,jkl->ikl', ke, u_selected) # compute x^penal * U.T @ ke @ U 70 | ce = anp.einsum('ijk,ijk->jk', u_selected, ke_u) 71 | C = young_modulus(x_phys, e_0, e_min, p=penal) * ce.T 72 | return anp.sum(C) 73 | 74 | def get_stiffness_matrix(e, nu): # e=young's modulus, nu=poisson coefficient 75 | k = anp.array([1/2-nu/6, 1/8+nu/8, -1/4-nu/12, -1/8+3*nu/8, 76 | -1/4+nu/12, -1/8-nu/8, nu/6, 1/8-3*nu/8]) 77 | return e/(1-nu**2)*anp.array([[k[0], k[1], k[2], k[3], k[4], k[5], k[6], k[7]], 78 | [k[1], k[0], k[7], k[6], k[5], k[4], k[3], k[2]], 79 | [k[2], k[7], k[0], k[5], k[6], k[3], k[4], k[1]], 80 | [k[3], k[6], k[5], k[0], k[7], k[2], k[1], k[4]], 81 | [k[4], k[5], k[6], k[7], k[0], k[1], k[2], k[3]], 82 | [k[5], k[4], k[3], k[2], k[1], k[0], k[7], k[6]], 83 | [k[6], k[3], k[4], k[1], k[2], k[7], k[0], k[5]], 84 | [k[7], k[2], k[1], k[4], k[3], k[6], k[5], k[0]]]) 85 | 86 | def get_k(stiffness, ke): 87 | # Constructs sparse stiffness matrix k (used in the displace fn) 88 | # First, get position of the nodes of each element in the stiffness matrix 89 | nely, nelx = stiffness.shape 90 | ely, elx = anp.meshgrid(range(nely), range(nelx)) # x, y coords 91 | ely, elx = ely.reshape(-1, 1), elx.reshape(-1, 1) 92 | 93 | n1 = (nely+1)*(elx+0) + (ely+0) 94 | n2 = (nely+1)*(elx+1) + (ely+0) 95 | n3 = (nely+1)*(elx+1) + (ely+1) 96 | n4 = (nely+1)*(elx+0) + (ely+1) 97 | edof = anp.array([2*n1, 2*n1+1, 2*n2, 2*n2+1, 2*n3, 2*n3+1, 2*n4, 2*n4+1]) 98 | edof = edof.T[0] 99 | x_list = anp.repeat(edof, 8) # flat list pointer of each node in an element 100 | y_list = anp.tile(edof, 8).flatten() # flat list pointer of each node in elem 101 | 102 | # make the global stiffness matrix K 103 | kd = stiffness.T.reshape(nelx*nely, 1, 1) 104 | value_list = (kd * anp.tile(ke, kd.shape)).flatten() 105 | return value_list, y_list, x_list 106 | 107 | def displace(x_phys, ke, forces, freedofs, fixdofs, *, penal=3, e_min=1e-9, e_0=1): 108 | # Displaces the load x using finite element techniques (solve_coo=most of runtime) 109 | stiffness = young_modulus(x_phys, e_0, e_min, p=penal) 110 | k_entries, k_ylist, k_xlist = get_k(stiffness, ke) 111 | 112 | index_map, keep, indices = _get_dof_indices(freedofs, fixdofs, k_ylist, k_xlist) 113 | 114 | u_nonzero = solve_coo(k_entries[keep], indices, forces[freedofs], sym_pos=True) 115 | u_values = anp.concatenate([u_nonzero, anp.zeros(len(fixdofs))]) 116 | return u_values[index_map] 117 | 118 | ##### Sparse matrix (COO) helper functions ##### 119 | def _get_dof_indices(freedofs, fixdofs, k_xlist, k_ylist): # sparse matrix helper function 120 | index_map = inverse_permutation(anp.concatenate([freedofs, fixdofs])) 121 | keep = anp.isin(k_xlist, freedofs) & anp.isin(k_ylist, freedofs) 122 | # Now we index an indexing array that is being indexed by the indices of k 123 | i = index_map[k_ylist][keep] 124 | j = index_map[k_xlist][keep] 125 | return index_map, keep, anp.stack([i, j]) 126 | 127 | def inverse_permutation(indices): # reverses an index operation 128 | inverse_perm = np.zeros(len(indices), dtype=anp.int64) 129 | inverse_perm[indices] = np.arange(len(indices), dtype=anp.int64) 130 | return inverse_perm 131 | 132 | def _get_solver(a_entries, a_indices, size, sym_pos): 133 | # a is (usu.) symmetric positive; could solve 2x faster w/sksparse.cholmod.cholesky(a).solve_A 134 | a = scipy.sparse.coo_matrix((a_entries, a_indices), shape=(size,)*2).tocsc() 135 | return scipy.sparse.linalg.splu(a).solve 136 | 137 | ##### Autograd custom gradients ##### 138 | @autograd.primitive 139 | def solve_coo(a_entries, a_indices, b, sym_pos=False): 140 | solver = _get_solver(a_entries, a_indices, b.size, sym_pos) 141 | return solver(b) 142 | 143 | def grad_solve_coo_entries(ans, a_entries, a_indices, b, sym_pos=False): 144 | def jvp(grad_ans): 145 | lambda_ = solve_coo(a_entries, a_indices if sym_pos else a_indices[::-1], 146 | grad_ans, sym_pos) 147 | i, j = a_indices 148 | return -lambda_[i] * ans[j] 149 | return jvp 150 | 151 | autograd.extend.defvjp(solve_coo, grad_solve_coo_entries, 152 | lambda: print('err: gradient undefined'), 153 | lambda: print('err: gradient not implemented')) 154 | 155 | @autograd.extend.primitive 156 | def gaussian_filter(x, width): # 2D gaussian blur/filter 157 | return scipy.ndimage.gaussian_filter(x, width, mode='reflect') 158 | 159 | def _gaussian_filter_vjp(ans, x, width): # gives the gradient of orig. function w.r.t. x 160 | del ans, x # unused 161 | return lambda g: gaussian_filter(g, width) 162 | autograd.extend.defvjp(gaussian_filter, _gaussian_filter_vjp) 163 | 164 | ##### Main optimization function ##### 165 | def fast_stopt(args, x=None, verbose=True): 166 | if x is None: 167 | x = anp.ones((args.nely, args.nelx)) * args.density # init mass 168 | 169 | reshape = lambda x: x.reshape(args.nely, args.nelx) 170 | objective_fn = lambda x: objective(reshape(x), args) # don't enforce mass constraint here 171 | constraint = lambda params: mean_density(reshape(params), args) - args.density 172 | 173 | def wrap_autograd_func(func, losses=None, frames=None): 174 | def wrapper(x, grad): 175 | if grad.size > 0: 176 | value, grad[:] = autograd.value_and_grad(func)(x) 177 | else: 178 | value = func(x) 179 | if losses is not None: 180 | losses.append(value) 181 | if frames is not None: 182 | frames.append(reshape(x).copy()) 183 | if verbose and len(frames) % args.print_every == 0: 184 | print('step {}, loss {:.2e}, t={:.2f}s'.format(len(frames), value, time.time()-dt)) 185 | return value 186 | return wrapper 187 | 188 | losses, frames = [], [] ; dt = time.time() 189 | print('Optimizing a problem with {} nodes'.format(len(args.forces))) 190 | opt = nlopt.opt(nlopt.LD_MMA, x.size) 191 | opt.set_lower_bounds(0.0) ; opt.set_upper_bounds(1.0) 192 | opt.set_min_objective(wrap_autograd_func(objective_fn, losses, frames)) 193 | opt.add_inequality_constraint(wrap_autograd_func(constraint), 1e-8) 194 | opt.set_maxeval(args.opt_steps + 1) 195 | opt.optimize(x.flatten()) 196 | return np.array(losses), reshape(frames[-1]), np.array(frames) 197 | 198 | ##### Run the simulation and visualize the result ##### 199 | args = get_args(*mbb_beam()) 200 | losses, x, mbb_frames = fast_stopt(args=args, verbose=True) 201 | plt.figure(dpi=100) ; print('\nFinal MBB beam design:') 202 | plt.imshow(np.concatenate([x[:,::-1],x], axis=1)) ; plt.show() -------------------------------------------------------------------------------- /stopt_240.py: -------------------------------------------------------------------------------- 1 | # A Tutorial on Structural Optimization | Sam Greydanus | 2022 2 | import time, functools, nlopt 3 | import numpy as np # for dense matrix ops 4 | import matplotlib.pyplot as plt # for plotting 5 | import autograd, autograd.core, autograd.extend, autograd.tracer # for adjoints 6 | import autograd.numpy as anp 7 | import scipy, scipy.ndimage, scipy.sparse, scipy.sparse.linalg # sparse matrices 8 | 9 | ##### Problem setup ##### 10 | class ObjectView(object): 11 | def __init__(self, d): self.__dict__ = d 12 | 13 | def get_args(normals, forces, density=0.4): # Manage the problem setup parameters 14 | width = normals.shape[0] - 1 15 | height = normals.shape[1] - 1 16 | fixdofs = np.flatnonzero(normals.ravel()) 17 | alldofs = np.arange(2 * (width + 1) * (height + 1)) 18 | freedofs = np.sort(list(set(alldofs) - set(fixdofs))) 19 | params = { 20 | # material properties 21 | 'young': 1, 'young_min': 1e-9, 'poisson': 0.3, 'g': 0, 22 | # constraints 23 | 'density': density, 'xmin': 0.001, 'xmax': 1.0, 24 | # input parameters 25 | 'nelx': width, 'nely': height, 'mask': 1, 'penal': 3.0, 'filter_width': 1, 26 | 'freedofs': freedofs, 'fixdofs': fixdofs, 'forces': forces.ravel(), 27 | # optimization parameters 28 | 'opt_steps': 80, 'print_every': 10} 29 | return ObjectView(params) 30 | 31 | def mbb_beam(width=80, height=25, density=0.4, y=1, x=0): # textbook beam example 32 | normals = np.zeros((width + 1, height + 1, 2)) 33 | normals[-1, -1, y] = 1 34 | normals[0, :, x] = 1 35 | forces = np.zeros((width + 1, height + 1, 2)) 36 | forces[0, 0, y] = -1 37 | return normals, forces, density 38 | 39 | def young_modulus(x, e_0, e_min, p=3): 40 | return e_min + x ** p * (e_0 - e_min) 41 | 42 | def physical_density(x, args, volume_contraint=False, use_filter=True): 43 | x = args.mask * x.reshape(args.nely, args.nelx) # reshape from 1D to 2D 44 | return gaussian_filter(x, args.filter_width) if use_filter else x # maybe filter 45 | 46 | def mean_density(x, args, volume_contraint=False, use_filter=True): 47 | return anp.mean(physical_density(x, args, volume_contraint, use_filter)) / anp.mean(args.mask) 48 | 49 | ##### Caching for NumPy arrays (gives a 2x speedup) ##### 50 | class _WrappedArray: 51 | """Hashable wrapper for NumPy arrays.""" 52 | def __init__(self, value): 53 | self.value = value 54 | 55 | def __eq__(self, other): 56 | return np.array_equal(self.value, other.value) 57 | 58 | def __hash__(self): 59 | # Collisions (there won't be many) mean __eq__ is checked: https://bit.ly/3EazdK3 60 | return hash(repr(self.value.ravel())) 61 | 62 | def ndarray_safe_lru_cache(maxsize=128): 63 | """An ndarray compatible version of functools.lru_cache.""" 64 | def decorator(func): 65 | @functools.lru_cache(maxsize) 66 | def cached_func(*args, **kwargs): 67 | args = tuple(a.value if isinstance(a, _WrappedArray) else a for a in args) 68 | kwargs = {k: v.value if isinstance(v, _WrappedArray) else v for k, v in kwargs.items()} 69 | return func(*args, **kwargs) 70 | 71 | @functools.wraps(func) 72 | def wrapper(*args, **kwargs): 73 | args = tuple(_WrappedArray(a) if isinstance(a, np.ndarray) else a for a in args) 74 | kwargs = {k: _WrappedArray(v) if isinstance(v, np.ndarray) else v for k, v in kwargs.items()} 75 | return cached_func(*args, **kwargs) 76 | return wrapper 77 | return decorator 78 | 79 | ##### Optimization objective + physics of elastic materials ##### 80 | def objective(x, args, volume_contraint=False, use_filter=True): 81 | kwargs = dict(penal=args.penal, e_min=args.young_min, e_0=args.young) 82 | x_phys = physical_density(x, args, volume_contraint=volume_contraint, use_filter=use_filter) 83 | ke = get_stiffness_matrix(args.young, args.poisson) # stiffness matrix 84 | u = displace(x_phys, ke, args.forces, args.freedofs, args.fixdofs, **kwargs) 85 | c = compliance(x_phys, u, ke, **kwargs) 86 | return c 87 | 88 | def compliance(x_phys, u, ke, *, penal=3, e_min=1e-9, e_0=1): 89 | nely, nelx = x_phys.shape 90 | ely, elx = anp.meshgrid(range(nely), range(nelx)) # x, y coords for the index map 91 | 92 | n1 = (nely+1)*(elx+0) + (ely+0) # nodes 93 | n2 = (nely+1)*(elx+1) + (ely+0) 94 | n3 = (nely+1)*(elx+1) + (ely+1) 95 | n4 = (nely+1)*(elx+0) + (ely+1) 96 | all_ixs = anp.array([2*n1, 2*n1+1, 2*n2, 2*n2+1, 2*n3, 2*n3+1, 2*n4, 2*n4+1]) 97 | u_selected = u[all_ixs] # select from u matrix 98 | 99 | ke_u = anp.einsum('ij,jkl->ikl', ke, u_selected) # compute x^penal * U.T @ ke @ U 100 | ce = anp.einsum('ijk,ijk->jk', u_selected, ke_u) 101 | C = young_modulus(x_phys, e_0, e_min, p=penal) * ce.T 102 | return anp.sum(C) 103 | 104 | def get_stiffness_matrix(e, nu): # e=young's modulus, nu=poisson coefficient 105 | k = anp.array([1/2-nu/6, 1/8+nu/8, -1/4-nu/12, -1/8+3*nu/8, 106 | -1/4+nu/12, -1/8-nu/8, nu/6, 1/8-3*nu/8]) 107 | return e/(1-nu**2)*anp.array([[k[0], k[1], k[2], k[3], k[4], k[5], k[6], k[7]], 108 | [k[1], k[0], k[7], k[6], k[5], k[4], k[3], k[2]], 109 | [k[2], k[7], k[0], k[5], k[6], k[3], k[4], k[1]], 110 | [k[3], k[6], k[5], k[0], k[7], k[2], k[1], k[4]], 111 | [k[4], k[5], k[6], k[7], k[0], k[1], k[2], k[3]], 112 | [k[5], k[4], k[3], k[2], k[1], k[0], k[7], k[6]], 113 | [k[6], k[3], k[4], k[1], k[2], k[7], k[0], k[5]], 114 | [k[7], k[2], k[1], k[4], k[3], k[6], k[5], k[0]]]) 115 | 116 | def get_k(stiffness, ke): 117 | # Constructs sparse stiffness matrix k (used in the displace fn) 118 | # First, get position of the nodes of each element in the stiffness matrix 119 | nely, nelx = stiffness.shape 120 | ely, elx = anp.meshgrid(range(nely), range(nelx)) # x, y coords 121 | ely, elx = ely.reshape(-1, 1), elx.reshape(-1, 1) 122 | 123 | n1 = (nely+1)*(elx+0) + (ely+0) 124 | n2 = (nely+1)*(elx+1) + (ely+0) 125 | n3 = (nely+1)*(elx+1) + (ely+1) 126 | n4 = (nely+1)*(elx+0) + (ely+1) 127 | edof = anp.array([2*n1, 2*n1+1, 2*n2, 2*n2+1, 2*n3, 2*n3+1, 2*n4, 2*n4+1]) 128 | edof = edof.T[0] 129 | x_list = anp.repeat(edof, 8) # flat list pointer of each node in an element 130 | y_list = anp.tile(edof, 8).flatten() # flat list pointer of each node in elem 131 | 132 | # make the global stiffness matrix K 133 | kd = stiffness.T.reshape(nelx*nely, 1, 1) 134 | value_list = (kd * anp.tile(ke, kd.shape)).flatten() 135 | return value_list, y_list, x_list 136 | 137 | def displace(x_phys, ke, forces, freedofs, fixdofs, *, penal=3, e_min=1e-9, e_0=1): 138 | # Displaces the load x using finite element techniques (solve_coo=most of runtime) 139 | stiffness = young_modulus(x_phys, e_0, e_min, p=penal) 140 | k_entries, k_ylist, k_xlist = get_k(stiffness, ke) 141 | index_map, keep, indices = _get_dof_indices(freedofs, fixdofs, k_ylist, k_xlist) 142 | 143 | u_nonzero = solve_coo(k_entries[keep], indices, forces[freedofs], sym_pos=True) 144 | u_values = anp.concatenate([u_nonzero, anp.zeros(len(fixdofs))]) 145 | return u_values[index_map] 146 | 147 | ##### Sparse matrix (COO) helper functions ##### 148 | @ndarray_safe_lru_cache(1) 149 | def _get_dof_indices(freedofs, fixdofs, k_xlist, k_ylist): # sparse matrix helper function 150 | index_map = inverse_permutation(anp.concatenate([freedofs, fixdofs])) 151 | keep = anp.isin(k_xlist, freedofs) & anp.isin(k_ylist, freedofs) 152 | # Now we index an indexing array that is being indexed by the indices of k 153 | i = index_map[k_ylist][keep] 154 | j = index_map[k_xlist][keep] 155 | return index_map, keep, anp.stack([i, j]) 156 | 157 | def inverse_permutation(indices): # reverses an index operation 158 | inverse_perm = np.zeros(len(indices), dtype=anp.int64) 159 | inverse_perm[indices] = np.arange(len(indices), dtype=anp.int64) 160 | return inverse_perm 161 | 162 | @ndarray_safe_lru_cache(1) 163 | def _get_solver(a_entries, a_indices, size, sym_pos): 164 | # a is (usu.) symmetric positive; could solve 2x faster w/sksparse.cholmod.cholesky(a).solve_A 165 | a = scipy.sparse.coo_matrix((a_entries, a_indices), shape=(size,)*2).tocsc() 166 | return scipy.sparse.linalg.splu(a).solve 167 | 168 | ##### Autograd custom gradients ##### 169 | @autograd.primitive 170 | def solve_coo(a_entries, a_indices, b, sym_pos=False): 171 | solver = _get_solver(a_entries, a_indices, b.size, sym_pos) 172 | return solver(b) 173 | 174 | def grad_solve_coo_entries(ans, a_entries, a_indices, b, sym_pos=False): 175 | def jvp(grad_ans): 176 | lambda_ = solve_coo(a_entries, a_indices if sym_pos else a_indices[::-1], 177 | grad_ans, sym_pos) 178 | i, j = a_indices 179 | return -lambda_[i] * ans[j] 180 | return jvp 181 | autograd.extend.defvjp(solve_coo, grad_solve_coo_entries, 182 | lambda: print('err: gradient undefined'), 183 | lambda: print('err: gradient not implemented')) 184 | 185 | @autograd.extend.primitive 186 | def gaussian_filter(x, width): # 2D gaussian blur/filter 187 | return scipy.ndimage.gaussian_filter(x, width, mode='reflect') 188 | 189 | def _gaussian_filter_vjp(ans, x, width): # gives the gradient of orig. function w.r.t. x 190 | del ans, x # unused 191 | return lambda g: gaussian_filter(g, width) 192 | autograd.extend.defvjp(gaussian_filter, _gaussian_filter_vjp) 193 | 194 | ##### Main optimization function ##### 195 | def fast_stopt(args, x=None, verbose=True): 196 | if x is None: 197 | x = anp.ones((args.nely, args.nelx)) * args.density # init mass 198 | reshape = lambda x: x.reshape(args.nely, args.nelx) 199 | objective_fn = lambda x: objective(reshape(x), args) # don't enforce mass constraint here 200 | constraint = lambda params: mean_density(reshape(params), args) - args.density 201 | 202 | def wrap_autograd_func(func, losses=None, frames=None): 203 | def wrapper(x, grad): 204 | if grad.size > 0: 205 | value, grad[:] = autograd.value_and_grad(func)(x) 206 | else: 207 | value = func(x) 208 | if losses is not None: 209 | losses.append(value) 210 | if frames is not None: 211 | frames.append(reshape(x).copy()) 212 | if verbose and len(frames) % args.print_every == 0: 213 | print('step {}, loss {:.2e}, t={:.2f}s'.format(len(frames), value, time.time()-dt)) 214 | return value 215 | return wrapper 216 | 217 | losses, frames = [], [] ; dt = time.time() 218 | print('Optimizing a problem with {} nodes'.format(len(args.forces))) 219 | opt = nlopt.opt(nlopt.LD_MMA, x.size) 220 | opt.set_lower_bounds(0.0) ; opt.set_upper_bounds(1.0) 221 | opt.set_min_objective(wrap_autograd_func(objective_fn, losses, frames)) 222 | opt.add_inequality_constraint(wrap_autograd_func(constraint), 1e-8) 223 | opt.set_maxeval(args.opt_steps + 1) 224 | opt.optimize(x.flatten()) 225 | return np.array(losses), reshape(frames[-1]), np.array(frames) 226 | 227 | ##### Run the simulation and visualize the result ##### 228 | args = get_args(*mbb_beam()) 229 | losses, x, mbb_frames = fast_stopt(args=args, verbose=True) 230 | plt.figure(dpi=100) ; print('\nFinal MBB beam design:') 231 | plt.imshow(np.concatenate([x[:,::-1],x], axis=1)) ; plt.show() --------------------------------------------------------------------------------