├── guided_diffusion ├── cis_list.py ├── __init__.py ├── __pycache__ │ ├── nn.cpython-311.pyc │ ├── unet.cpython-311.pyc │ ├── logger.cpython-311.pyc │ ├── losses.cpython-311.pyc │ ├── respace.cpython-311.pyc │ ├── __init__.cpython-311.pyc │ ├── dist_util.cpython-311.pyc │ ├── fp16_util.cpython-311.pyc │ ├── resample.cpython-311.pyc │ ├── simulation.cpython-311.pyc │ ├── train_util.cpython-311.pyc │ ├── script_util.cpython-311.pyc │ ├── image_datasets.cpython-311.pyc │ └── gaussian_diffusion.cpython-311.pyc ├── README.md ├── dist_util.py ├── losses.py ├── respace.py ├── nn.py ├── resample.py ├── image_datasets.py ├── fp16_util.py ├── train_util.py ├── script_util.py ├── logger.py └── simulation.py ├── images ├── logo.png ├── .DS_Store ├── Result1-2.png ├── Result1.png ├── Result2-2.png ├── Result2.png ├── result2-2.png ├── comparison.png ├── Result-waveguide.png └── Result-colorrouter.png ├── requirements.txt ├── 01-train.sh ├── LICENSE ├── 01-sample.sh ├── training_examples.ipynb ├── 02-train.ipynb ├── dataset_generation.py ├── sampling_example.ipynb ├── image_train.py ├── 02-sample.ipynb ├── image_sample.py ├── README.md ├── baseline_algorithms ├── [adjoint_optimization] GA_waveguide.py ├── [adjoint_optimization] nlopt_waveguide.py ├── [adjoint_optimization] GA_CIS.py └── [adjoint_optimization] nlopt_CIS.py └── simulation.py /guided_diffusion/cis_list.py: -------------------------------------------------------------------------------- 1 | red_list = [] 2 | blue_list = [] 3 | green_list = [] -------------------------------------------------------------------------------- /images/logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dongjin-seo2020/AdjointDiffusion/HEAD/images/logo.png -------------------------------------------------------------------------------- /images/.DS_Store: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dongjin-seo2020/AdjointDiffusion/HEAD/images/.DS_Store -------------------------------------------------------------------------------- /images/Result1-2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dongjin-seo2020/AdjointDiffusion/HEAD/images/Result1-2.png -------------------------------------------------------------------------------- /images/Result1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dongjin-seo2020/AdjointDiffusion/HEAD/images/Result1.png -------------------------------------------------------------------------------- /images/Result2-2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dongjin-seo2020/AdjointDiffusion/HEAD/images/Result2-2.png -------------------------------------------------------------------------------- /images/Result2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dongjin-seo2020/AdjointDiffusion/HEAD/images/Result2.png -------------------------------------------------------------------------------- /images/result2-2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dongjin-seo2020/AdjointDiffusion/HEAD/images/result2-2.png -------------------------------------------------------------------------------- /images/comparison.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dongjin-seo2020/AdjointDiffusion/HEAD/images/comparison.png -------------------------------------------------------------------------------- /guided_diffusion/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Codebase for "Improved Denoising Diffusion Probabilistic Models". 3 | """ 4 | -------------------------------------------------------------------------------- /images/Result-waveguide.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dongjin-seo2020/AdjointDiffusion/HEAD/images/Result-waveguide.png -------------------------------------------------------------------------------- /images/Result-colorrouter.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dongjin-seo2020/AdjointDiffusion/HEAD/images/Result-colorrouter.png -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/nn.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dongjin-seo2020/AdjointDiffusion/HEAD/guided_diffusion/__pycache__/nn.cpython-311.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/unet.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dongjin-seo2020/AdjointDiffusion/HEAD/guided_diffusion/__pycache__/unet.cpython-311.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/logger.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dongjin-seo2020/AdjointDiffusion/HEAD/guided_diffusion/__pycache__/logger.cpython-311.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/losses.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dongjin-seo2020/AdjointDiffusion/HEAD/guided_diffusion/__pycache__/losses.cpython-311.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/respace.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dongjin-seo2020/AdjointDiffusion/HEAD/guided_diffusion/__pycache__/respace.cpython-311.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/__init__.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dongjin-seo2020/AdjointDiffusion/HEAD/guided_diffusion/__pycache__/__init__.cpython-311.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/dist_util.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dongjin-seo2020/AdjointDiffusion/HEAD/guided_diffusion/__pycache__/dist_util.cpython-311.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/fp16_util.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dongjin-seo2020/AdjointDiffusion/HEAD/guided_diffusion/__pycache__/fp16_util.cpython-311.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/resample.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dongjin-seo2020/AdjointDiffusion/HEAD/guided_diffusion/__pycache__/resample.cpython-311.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/simulation.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dongjin-seo2020/AdjointDiffusion/HEAD/guided_diffusion/__pycache__/simulation.cpython-311.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/train_util.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dongjin-seo2020/AdjointDiffusion/HEAD/guided_diffusion/__pycache__/train_util.cpython-311.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/script_util.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dongjin-seo2020/AdjointDiffusion/HEAD/guided_diffusion/__pycache__/script_util.cpython-311.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/image_datasets.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dongjin-seo2020/AdjointDiffusion/HEAD/guided_diffusion/__pycache__/image_datasets.cpython-311.pyc -------------------------------------------------------------------------------- /guided_diffusion/__pycache__/gaussian_diffusion.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dongjin-seo2020/AdjointDiffusion/HEAD/guided_diffusion/__pycache__/gaussian_diffusion.cpython-311.pyc -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | # Note that the meep library has to be installed separately for the compatibility of other libraries, such as torch. For further information, refer to the README. 2 | 3 | numpy==1.26.4 4 | scipy>=1.11 5 | matplotlib>=3.8 6 | pillow>=10.0 7 | tqdm>=4.66 8 | wandb>=0.16 9 | nlopt>=2.7.1 10 | autograd>=1.6 11 | scikit-image>=0.22 12 | blobfile==2.1.1 13 | mpi4py==3.1.6 -------------------------------------------------------------------------------- /guided_diffusion/README.md: -------------------------------------------------------------------------------- 1 | # guided_diffusion 2 | 3 | This directory contains code adapted from [OpenAI's guided-diffusion](https://github.com/openai/guided-diffusion), originally licensed under the MIT License. 4 | 5 | All modifications made in this repository are for the purpose of integrating physics-guided optimization via adjoint gradients. 6 | 7 | The original license is preserved in the root `LICENSE` file. -------------------------------------------------------------------------------- /01-train.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | DATA_DIR=... 4 | LOG_DIR=... 5 | GPU_ID=... 6 | 7 | # Set model, diffusion, directory, and training flags 8 | MODEL_FLAGS="--dropout 0.1 --class_cond True --gray_imgs True" 9 | DIFF_FLAGS="--learn_sigma True --diffusion_steps 1000 --noise_schedule cosine" 10 | DIR_FLAGS="--data_dir ${DATA_DIR} --log_dir ${LOG_DIR}" 11 | TRAIN_FLAGS="--lr 1e-4 --batch_size 128 --microbatch 64 --gpu_id ${GPU_ID}" 12 | 13 | # Run training 14 | python3 image_train.py $MODEL_FLAGS $DIFF_FLAGS $DIR_FLAGS $TRAIN_FLAGS 15 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 OpenAI 4 | Copyright (c) 2025 Dongjin Seo 5 | 6 | Permission is hereby granted, free of charge, to any person obtaining a copy 7 | of this software and associated documentation files (the "Software"), to deal 8 | in the Software without restriction, including without limitation the rights 9 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 10 | copies of the Software, and to permit persons to whom the Software is 11 | furnished to do so, subject to the following conditions: 12 | 13 | The above copyright notice and this permission notice shall be included in all 14 | copies or substantial portions of the Software. 15 | 16 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 17 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 18 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 19 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 20 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 21 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 22 | SOFTWARE. -------------------------------------------------------------------------------- /01-sample.sh: -------------------------------------------------------------------------------- 1 | gpu_id=0 2 | prop="top" 3 | MODEL_PATH=... 4 | LOG_PATH=... 5 | 6 | for manual_class_id in 0 7 | do 8 | for tsr in 100 9 | do 10 | eta_list=(1) 11 | for eta in "${eta_list[@]}" 12 | do 13 | MODEL_FLAGS="--dropout 0.1 --class_cond True --gray_imgs True" 14 | DIFF_FLAGS="--learn_sigma True --diffusion_steps 1000 --noise_schedule cosine" 15 | DIR_FLAGS="--model_path ${MODEL_PATH} \ 16 | --log_dir ${LOG_PATH}/sim-guided/${prop}_tsr=${tsr}_class=${manual_class_id}_eta=${eta}" 17 | SAMPLE_FLAGS="--batch_size 1 --num_samples 1 --timestep_respacing ${tsr} --num_classes 3 --manual_class_id ${manual_class_id} --gpu_id ${gpu_id} --save_img False" 18 | SIM_FLAGS="--sim_guided True --sim_type waveguide --use_normed_grad True --use_adjgrad_norm False --eta ${eta} --prop_dir ${prop} --save_inter True --interval 1" 19 | echo -e "\n\n\n\n############################ Sampling with eta = ${eta} ##############################\n" 20 | python3 image_sample.py $MODEL_FLAGS $DIFF_FLAGS $DIR_FLAGS $SAMPLE_FLAGS $SIM_FLAGS 21 | done 22 | done 23 | done 24 | -------------------------------------------------------------------------------- /training_examples.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": null, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "MODEL_FLAGS=\"--dropout 0.1 --class_cond True --gray_imgs True\"\n", 10 | "DIFF_FLAGS=\"--learn_sigma True --diffusion_steps 1000 --noise_schedule cosine\"\n", 11 | "DIR_FLAGS=\"--data_dir /media/usb_media/datasets/physics/wg/64/class-cond --log_dir /media/usb_media/guided-diffusion-phy/wg-64_class-cond\"\n", 12 | "TRAIN_FLAGS=\"--lr 1e-4 --batch_size 128 --microbatch 64 --gpu_id 7\"\n", 13 | "\n", 14 | "!python3 image_train.py $MODEL_FLAGS $DIFF_FLAGS $DIR_FLAGS $TRAIN_FLAGS" 15 | ] 16 | } 17 | ], 18 | "metadata": { 19 | "kernelspec": { 20 | "display_name": "physics-diff", 21 | "language": "python", 22 | "name": "python3" 23 | }, 24 | "language_info": { 25 | "codemirror_mode": { 26 | "name": "ipython", 27 | "version": 3 28 | }, 29 | "file_extension": ".py", 30 | "mimetype": "text/x-python", 31 | "name": "python", 32 | "nbconvert_exporter": "python", 33 | "pygments_lexer": "ipython3", 34 | "version": "3.11.7" 35 | } 36 | }, 37 | "nbformat": 4, 38 | "nbformat_minor": 2 39 | } 40 | -------------------------------------------------------------------------------- /02-train.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": null, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "DATA_DIR = \"/media/usb_media/datasets/physics/wg/64/class-cond\"\n", 10 | "LOG_DIR = \"/media/usb_media/guided-diffusion-phy/wg-64_class-cond\"\n", 11 | "GPU_ID = 7\n", 12 | "\n", 13 | "MODEL_FLAGS = \"--dropout 0.1 --class_cond True --gray_imgs True\"\n", 14 | "DIFF_FLAGS = \"--learn_sigma True --diffusion_steps 1000 --noise_schedule cosine\"\n", 15 | "DIR_FLAGS = f\"--data_dir {DATA_DIR} --log_dir {LOG_DIR}\"\n", 16 | "TRAIN_FLAGS = f\"--lr 1e-4 --batch_size 128 --microbatch 64 --gpu_id {GPU_ID}\"\n", 17 | "\n", 18 | "!python3 image_train.py $MODEL_FLAGS $DIFF_FLAGS $DIR_FLAGS $TRAIN_FLAGS" 19 | ] 20 | } 21 | ], 22 | "metadata": { 23 | "kernelspec": { 24 | "display_name": "physics-diff", 25 | "language": "python", 26 | "name": "python3" 27 | }, 28 | "language_info": { 29 | "codemirror_mode": { 30 | "name": "ipython", 31 | "version": 3 32 | }, 33 | "file_extension": ".py", 34 | "mimetype": "text/x-python", 35 | "name": "python", 36 | "nbconvert_exporter": "python", 37 | "pygments_lexer": "ipython3", 38 | "version": "3.11.7" 39 | } 40 | }, 41 | "nbformat": 4, 42 | "nbformat_minor": 2 43 | } 44 | -------------------------------------------------------------------------------- /dataset_generation.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import matplotlib.pyplot as plt 3 | from scipy.ndimage import gaussian_filter 4 | from PIL import Image 5 | 6 | import os 7 | 8 | def generate_random_binary_structure(n): 9 | """Generate a random binary structure of size n x n.""" 10 | return np.random.randint(0, 2, (n, n), dtype=np.uint8)*255 11 | 12 | def apply_gaussian_filter(image, sigma): 13 | """Apply a Gaussian filter to an image.""" 14 | return gaussian_filter(image, sigma=sigma) 15 | 16 | def binarize_image(image, threshold=255/2): 17 | """Binarize an image based on a threshold.""" 18 | return (image > threshold).astype(int)*255 19 | 20 | def save_image_as_png(array, filename='image.png'): 21 | """Save the image as a PNG.""" 22 | #image = Image.fromarray(array,"L") 23 | #image.save(filename) 24 | plt.imsave(filename, array, cmap='gray') 25 | 26 | k = 2 27 | n = 64 28 | 29 | datasize = 30000 30 | 31 | 32 | directory = 'datasets/'+str(n)+'/sigma'+str(k)+'/struct/' 33 | 34 | import os 35 | if not os.path.exists(directory): 36 | os.makedirs(directory) 37 | 38 | 39 | for i in range(datasize): 40 | binary_structure = generate_random_binary_structure(n) 41 | noisy_structure = apply_gaussian_filter(binary_structure, k) 42 | binarized_structure = binarize_image(noisy_structure) 43 | print(binarized_structure) 44 | save_image_as_png(binarized_structure, filename='datasets/'+str(n)+'/sigma'+str(k)+'/struct/'+str(i)+'.png') 45 | 46 | 47 | 48 | with Image.open('datasets/'+str(n)+'/sigma'+str(k)+'/struct/1.png') as img: 49 | width, height = img.size 50 | 51 | print(width, height) 52 | print(img.mode) -------------------------------------------------------------------------------- /sampling_example.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "metadata": {}, 6 | "source": [ 7 | "# simulation-guided" 8 | ] 9 | }, 10 | { 11 | "cell_type": "code", 12 | "execution_count": null, 13 | "metadata": {}, 14 | "outputs": [], 15 | "source": [ 16 | "import numpy as np\n", 17 | "\n", 18 | "for manual_class_id in [0]:\n", 19 | " for tsr in [250]:\n", 20 | " for eta in [2.0, 4.0, 8.0, 12.0, 16.0, 20.0]:\n", 21 | " MODEL_FLAGS=\"--dropout 0.1 --class_cond True --gray_imgs True\"\n", 22 | " DIFF_FLAGS=\"--learn_sigma True --diffusion_steps 1000 --noise_schedule cosine\"\n", 23 | " DIR_FLAGS=\"--model_path /media/usb_media/guided-diffusion-phy/wg-64_class-cond/ema_0.9999_025000.pt \" \\\n", 24 | " f\"--log_dir /media/usb_media/guided-diffusion-phy/wg-64_class-cond/imgs/model_025000/sim-guided/ltt_tsr={tsr}_class={manual_class_id}_eta={eta}\"\n", 25 | " SAMPLE_FLAGS=f\"--batch_size 1 --num_samples 1 --timestep_respacing {tsr} --num_classes 3 --manual_class_id {manual_class_id} --gpu_id 4 --save_img False\"\n", 26 | " SIM_FLAGS=f\"--sim_guided True --eta {eta} --use_normed_grad True --save_inter True --interval 5\"\n", 27 | " print(\"\\n\\n\\n\\n############################ Sampling with eta = \", eta, \" ############################\\n\")\n", 28 | " !python image_sample.py $MODEL_FLAGS $DIFF_FLAGS $DIR_FLAGS $SAMPLE_FLAGS $SIM_FLAGS" 29 | ] 30 | } 31 | ], 32 | "metadata": { 33 | "kernelspec": { 34 | "display_name": "physics-diff", 35 | "language": "python", 36 | "name": "python3" 37 | }, 38 | "language_info": { 39 | "codemirror_mode": { 40 | "name": "ipython", 41 | "version": 3 42 | }, 43 | "file_extension": ".py", 44 | "mimetype": "text/x-python", 45 | "name": "python", 46 | "nbconvert_exporter": "python", 47 | "pygments_lexer": "ipython3", 48 | "version": "3.11.7" 49 | } 50 | }, 51 | "nbformat": 4, 52 | "nbformat_minor": 2 53 | } 54 | -------------------------------------------------------------------------------- /guided_diffusion/dist_util.py: -------------------------------------------------------------------------------- 1 | """ 2 | Helpers for distributed training. 3 | """ 4 | 5 | import io 6 | import os 7 | import socket 8 | 9 | import blobfile as bf 10 | from mpi4py import MPI 11 | import torch as th 12 | import torch.distributed as dist 13 | 14 | # Change this to reflect your cluster layout. 15 | # The GPU for a given rank is (rank % GPUS_PER_NODE). 16 | GPUS_PER_NODE = 8 17 | 18 | SETUP_RETRY_COUNT = 3 19 | 20 | 21 | def setup_dist(): 22 | """ 23 | Setup a distributed process group. 24 | """ 25 | if dist.is_initialized(): 26 | return 27 | os.environ["CUDA_VISIBLE_DEVICES"] = f"{MPI.COMM_WORLD.Get_rank() % GPUS_PER_NODE}" 28 | 29 | comm = MPI.COMM_WORLD 30 | backend = "gloo" if not th.cuda.is_available() else "nccl" 31 | 32 | if backend == "gloo": 33 | hostname = "localhost" 34 | else: 35 | hostname = socket.gethostbyname(socket.getfqdn()) 36 | os.environ["MASTER_ADDR"] = comm.bcast(hostname, root=0) 37 | os.environ["RANK"] = str(comm.rank) 38 | os.environ["WORLD_SIZE"] = str(comm.size) 39 | 40 | port = comm.bcast(_find_free_port(), root=0) 41 | os.environ["MASTER_PORT"] = str(port) 42 | dist.init_process_group(backend=backend, init_method="env://") 43 | 44 | 45 | def dev(): 46 | """ 47 | Get the device to use for torch.distributed. 48 | """ 49 | if th.cuda.is_available(): 50 | return th.device(f"cuda") 51 | return th.device("cpu") 52 | 53 | 54 | def load_state_dict(path, **kwargs): 55 | """ 56 | Load a PyTorch file without redundant fetches across MPI ranks. 57 | """ 58 | chunk_size = 2 ** 30 # MPI has a relatively small size limit 59 | if MPI.COMM_WORLD.Get_rank() == 0: 60 | with bf.BlobFile(path, "rb") as f: 61 | data = f.read() 62 | num_chunks = len(data) // chunk_size 63 | if len(data) % chunk_size: 64 | num_chunks += 1 65 | MPI.COMM_WORLD.bcast(num_chunks) 66 | for i in range(0, len(data), chunk_size): 67 | MPI.COMM_WORLD.bcast(data[i : i + chunk_size]) 68 | else: 69 | num_chunks = MPI.COMM_WORLD.bcast(None) 70 | data = bytes() 71 | for _ in range(num_chunks): 72 | data += MPI.COMM_WORLD.bcast(None) 73 | 74 | return th.load(io.BytesIO(data), **kwargs) 75 | 76 | 77 | def sync_params(params): 78 | """ 79 | Synchronize a sequence of Tensors across ranks from rank 0. 80 | """ 81 | for p in params: 82 | with th.no_grad(): 83 | dist.broadcast(p, 0) 84 | 85 | 86 | def _find_free_port(): 87 | try: 88 | s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 89 | s.bind(("", 0)) 90 | s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) 91 | return s.getsockname()[1] 92 | finally: 93 | s.close() 94 | -------------------------------------------------------------------------------- /guided_diffusion/losses.py: -------------------------------------------------------------------------------- 1 | """ 2 | Helpers for various likelihood-based losses. These are ported from the original 3 | Ho et al. diffusion models codebase: 4 | https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/utils.py 5 | """ 6 | 7 | import numpy as np 8 | 9 | import torch as th 10 | 11 | 12 | def normal_kl(mean1, logvar1, mean2, logvar2): 13 | """ 14 | Compute the KL divergence between two gaussians. 15 | 16 | Shapes are automatically broadcasted, so batches can be compared to 17 | scalars, among other use cases. 18 | """ 19 | tensor = None 20 | for obj in (mean1, logvar1, mean2, logvar2): 21 | if isinstance(obj, th.Tensor): 22 | tensor = obj 23 | break 24 | assert tensor is not None, "at least one argument must be a Tensor" 25 | 26 | # Force variances to be Tensors. Broadcasting helps convert scalars to 27 | # Tensors, but it does not work for th.exp(). 28 | logvar1, logvar2 = [ 29 | x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor) 30 | for x in (logvar1, logvar2) 31 | ] 32 | 33 | return 0.5 * ( 34 | -1.0 35 | + logvar2 36 | - logvar1 37 | + th.exp(logvar1 - logvar2) 38 | + ((mean1 - mean2) ** 2) * th.exp(-logvar2) 39 | ) 40 | 41 | 42 | def approx_standard_normal_cdf(x): 43 | """ 44 | A fast approximation of the cumulative distribution function of the 45 | standard normal. 46 | """ 47 | return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3)))) 48 | 49 | 50 | def discretized_gaussian_log_likelihood(x, *, means, log_scales): 51 | """ 52 | Compute the log-likelihood of a Gaussian distribution discretizing to a 53 | given image. 54 | 55 | :param x: the target images. It is assumed that this was uint8 values, 56 | rescaled to the range [-1, 1]. 57 | :param means: the Gaussian mean Tensor. 58 | :param log_scales: the Gaussian log stddev Tensor. 59 | :return: a tensor like x of log probabilities (in nats). 60 | """ 61 | assert x.shape == means.shape == log_scales.shape 62 | centered_x = x - means 63 | inv_stdv = th.exp(-log_scales) 64 | plus_in = inv_stdv * (centered_x + 1.0 / 255.0) 65 | cdf_plus = approx_standard_normal_cdf(plus_in) 66 | min_in = inv_stdv * (centered_x - 1.0 / 255.0) 67 | cdf_min = approx_standard_normal_cdf(min_in) 68 | log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12)) 69 | log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12)) 70 | cdf_delta = cdf_plus - cdf_min 71 | log_probs = th.where( 72 | x < -0.999, 73 | log_cdf_plus, 74 | th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))), 75 | ) 76 | assert log_probs.shape == x.shape 77 | return log_probs 78 | -------------------------------------------------------------------------------- /image_train.py: -------------------------------------------------------------------------------- 1 | """ 2 | Train a diffusion model on images. 3 | """ 4 | 5 | import argparse 6 | import torch as th 7 | 8 | from guided_diffusion import dist_util, logger 9 | from guided_diffusion.image_datasets import load_data 10 | from guided_diffusion.resample import create_named_schedule_sampler 11 | from guided_diffusion.script_util import ( 12 | model_and_diffusion_defaults, 13 | create_model_and_diffusion, 14 | args_to_dict, 15 | add_dict_to_argparser, 16 | ) 17 | from guided_diffusion.train_util import TrainLoop 18 | 19 | 20 | def main(): 21 | args = create_argparser().parse_args() 22 | logger.configure(dir=args.log_dir) 23 | 24 | if not args.gpu_id == '': 25 | logger.log("using device %s" % args.gpu_id) 26 | th.cuda.set_device(th.device(f"cuda:{int(args.gpu_id)}")) 27 | 28 | dist_util.setup_dist() 29 | 30 | logger.log("creating model and diffusion...") 31 | model, diffusion = create_model_and_diffusion( 32 | **args_to_dict(args, model_and_diffusion_defaults().keys()) 33 | ) 34 | model.to(dist_util.dev()) 35 | schedule_sampler = create_named_schedule_sampler(args.schedule_sampler, diffusion) 36 | 37 | logger.log("creating data loader...") 38 | data = load_data( 39 | data_dir=args.data_dir, 40 | batch_size=args.batch_size, 41 | image_size=args.image_size, 42 | class_cond=args.class_cond, 43 | gray_imgs=args.gray_imgs, ##### modified by sum: for admitting grayscale images ##### 44 | ) 45 | 46 | logger.log("training...") 47 | TrainLoop( 48 | model=model, 49 | diffusion=diffusion, 50 | data=data, 51 | batch_size=args.batch_size, 52 | microbatch=args.microbatch, 53 | lr=args.lr, 54 | ema_rate=args.ema_rate, 55 | log_interval=args.log_interval, 56 | save_interval=args.save_interval, 57 | resume_checkpoint=args.resume_checkpoint, 58 | use_fp16=args.use_fp16, 59 | fp16_scale_growth=args.fp16_scale_growth, 60 | schedule_sampler=schedule_sampler, 61 | weight_decay=args.weight_decay, 62 | lr_anneal_steps=args.lr_anneal_steps, 63 | ).run_loop() 64 | 65 | 66 | def create_argparser(): 67 | defaults = dict( 68 | data_dir="datasets", 69 | log_dir="./out_dir", 70 | gpu_id='', 71 | schedule_sampler="uniform", 72 | lr=1e-4, 73 | weight_decay=0.0, 74 | lr_anneal_steps=0, 75 | batch_size=1, 76 | microbatch=-1, # -1 disables microbatches 77 | ema_rate="0.9999", # comma-separated list of EMA values 78 | log_interval=100, 79 | save_interval=1000, 80 | resume_checkpoint="", 81 | use_fp16=False, 82 | fp16_scale_growth=1e-3, 83 | gray_imgs=True, ##### modified by sum: for admitting grayscale images ##### 84 | ) 85 | defaults.update(model_and_diffusion_defaults()) 86 | parser = argparse.ArgumentParser() 87 | add_dict_to_argparser(parser, defaults) 88 | return parser 89 | 90 | 91 | if __name__ == "__main__": 92 | main() 93 | -------------------------------------------------------------------------------- /02-sample.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": null, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "MODEL_PATH = \"/media/usb_media/guided-diffusion-phy/wg-64_class-cond/ema_0.9999_025000.pt\"\n", 10 | "LOG_DIR = \"/media/usb_media/guided-diffusion-phy/wg-64_class-cond/imgs/model_025000\"\n", 11 | "GPU_ID = ...\n", 12 | "\n", 13 | "for tsr in [1000, 500, 250, 100]:\n", 14 | " for manual_class_id in range(3):\n", 15 | "\n", 16 | " MODEL_FLAGS = \"--dropout 0.1 --class_cond True --gray_imgs True\"\n", 17 | " DIFF_FLAGS = \"--learn_sigma True --diffusion_steps 1000 --noise_schedule cosine\"\n", 18 | " DIR_FLAGS = f\"--model_path {MODEL_PATH} \" \\\n", 19 | " f\"--log_dir {LOG_DIR}/tsr={tsr}/class={manual_class_id}\"\n", 20 | " SAMPLE_FLAGS = f\"--batch_size 8 --num_samples 8 --timestep_respacing {tsr} --num_classes 3 --manual_class_id {manual_class_id} --gpu_id {GPU_ID} --save_img True\"\n", 21 | "\n", 22 | " !python image_sample.py $MODEL_FLAGS $DIFF_FLAGS $DIR_FLAGS $SAMPLE_FLAGS" 23 | ] 24 | }, 25 | { 26 | "cell_type": "markdown", 27 | "metadata": {}, 28 | "source": [ 29 | "# simulation-guided" 30 | ] 31 | }, 32 | { 33 | "cell_type": "code", 34 | "execution_count": null, 35 | "metadata": {}, 36 | "outputs": [], 37 | "source": [ 38 | "import numpy as np\n", 39 | "\n", 40 | "for manual_class_id in [0, 1, 2]:\n", 41 | " for tsr in [5, 10, 20]:\n", 42 | " for eta in [2.0, 4.0, 8.0, 12.0, 16.0, 20.0]:\n", 43 | " MODEL_FLAGS = \"--dropout 0.1 --class_cond True --gray_imgs True\"\n", 44 | " DIFF_FLAGS = \"--learn_sigma True --diffusion_steps 1000 --noise_schedule cosine\"\n", 45 | " DIR_FLAGS = \"--model_path /media/usb_media/guided-diffusion-phy/wg-64_class-cond/ema_0.9999_025000.pt \" \\\n", 46 | " f\"--log_dir /media/usb_media/guided-diffusion-phy/wg-64_class-cond/imgs/model_025000/sim-guided/ltt_tsr={tsr}_class={manual_class_id}_eta={eta}\"\n", 47 | " SAMPLE_FLAGS = f\"--batch_size 1 --num_samples 1 --timestep_respacing {tsr} --num_classes 3 --manual_class_id {manual_class_id} --gpu_id 6 --save_img False\"\n", 48 | " SIM_FLAGS = f\"--sim_guided True --eta {eta} --use_normed_grad True --save_inter True --interval 5\"\n", 49 | "\n", 50 | " !python image_sample.py $MODEL_FLAGS $DIFF_FLAGS $DIR_FLAGS $SAMPLE_FLAGS $SIM_FLAGS" 51 | ] 52 | } 53 | ], 54 | "metadata": { 55 | "kernelspec": { 56 | "display_name": "physics-diff", 57 | "language": "python", 58 | "name": "python3" 59 | }, 60 | "language_info": { 61 | "codemirror_mode": { 62 | "name": "ipython", 63 | "version": 3 64 | }, 65 | "file_extension": ".py", 66 | "mimetype": "text/x-python", 67 | "name": "python", 68 | "nbconvert_exporter": "python", 69 | "pygments_lexer": "ipython3", 70 | "version": "3.11.7" 71 | } 72 | }, 73 | "nbformat": 4, 74 | "nbformat_minor": 2 75 | } 76 | -------------------------------------------------------------------------------- /guided_diffusion/respace.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch as th 3 | 4 | from .gaussian_diffusion import GaussianDiffusion 5 | 6 | 7 | def space_timesteps(num_timesteps, section_counts): 8 | """ 9 | Create a list of timesteps to use from an original diffusion process, 10 | given the number of timesteps we want to take from equally-sized portions 11 | of the original process. 12 | 13 | For example, if there's 300 timesteps and the section counts are [10,15,20] 14 | then the first 100 timesteps are strided to be 10 timesteps, the second 100 15 | are strided to be 15 timesteps, and the final 100 are strided to be 20. 16 | 17 | If the stride is a string starting with "ddim", then the fixed striding 18 | from the DDIM paper is used, and only one section is allowed. 19 | 20 | :param num_timesteps: the number of diffusion steps in the original 21 | process to divide up. 22 | :param section_counts: either a list of numbers, or a string containing 23 | comma-separated numbers, indicating the step count 24 | per section. As a special case, use "ddimN" where N 25 | is a number of steps to use the striding from the 26 | DDIM paper. 27 | :return: a set of diffusion steps from the original process to use. 28 | """ 29 | if isinstance(section_counts, str): 30 | if section_counts.startswith("ddim"): 31 | desired_count = int(section_counts[len("ddim") :]) 32 | for i in range(1, num_timesteps): 33 | if len(range(0, num_timesteps, i)) == desired_count: 34 | return set(range(0, num_timesteps, i)) 35 | raise ValueError( 36 | f"cannot create exactly {num_timesteps} steps with an integer stride" 37 | ) 38 | section_counts = [int(x) for x in section_counts.split(",")] 39 | size_per = num_timesteps // len(section_counts) 40 | extra = num_timesteps % len(section_counts) 41 | start_idx = 0 42 | all_steps = [] 43 | for i, section_count in enumerate(section_counts): 44 | size = size_per + (1 if i < extra else 0) 45 | if size < section_count: 46 | raise ValueError( 47 | f"cannot divide section of {size} steps into {section_count}" 48 | ) 49 | if section_count <= 1: 50 | frac_stride = 1 51 | else: 52 | frac_stride = (size - 1) / (section_count - 1) 53 | cur_idx = 0.0 54 | taken_steps = [] 55 | for _ in range(section_count): 56 | taken_steps.append(start_idx + round(cur_idx)) 57 | cur_idx += frac_stride 58 | all_steps += taken_steps 59 | start_idx += size 60 | return set(all_steps) 61 | 62 | 63 | class SpacedDiffusion(GaussianDiffusion): 64 | """ 65 | A diffusion process which can skip steps in a base diffusion process. 66 | 67 | :param use_timesteps: a collection (sequence or set) of timesteps from the 68 | original diffusion process to retain. 69 | :param kwargs: the kwargs to create the base diffusion process. 70 | """ 71 | 72 | def __init__(self, use_timesteps, **kwargs): 73 | self.use_timesteps = set(use_timesteps) 74 | self.timestep_map = [] 75 | self.original_num_steps = len(kwargs["betas"]) 76 | 77 | base_diffusion = GaussianDiffusion(**kwargs) # pylint: disable=missing-kwoa 78 | last_alpha_cumprod = 1.0 79 | new_betas = [] 80 | for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod): 81 | if i in self.use_timesteps: 82 | new_betas.append(1 - alpha_cumprod / last_alpha_cumprod) 83 | last_alpha_cumprod = alpha_cumprod 84 | self.timestep_map.append(i) 85 | kwargs["betas"] = np.array(new_betas) 86 | super().__init__(**kwargs) 87 | 88 | def p_mean_variance( 89 | self, model, *args, **kwargs 90 | ): # pylint: disable=signature-differs 91 | return super().p_mean_variance(self._wrap_model(model), *args, **kwargs) 92 | 93 | def training_losses( 94 | self, model, *args, **kwargs 95 | ): # pylint: disable=signature-differs 96 | return super().training_losses(self._wrap_model(model), *args, **kwargs) 97 | 98 | def condition_mean(self, cond_fn, *args, **kwargs): 99 | return super().condition_mean(self._wrap_model(cond_fn), *args, **kwargs) 100 | 101 | def condition_score(self, cond_fn, *args, **kwargs): 102 | return super().condition_score(self._wrap_model(cond_fn), *args, **kwargs) 103 | 104 | def _wrap_model(self, model): 105 | if isinstance(model, _WrappedModel): 106 | return model 107 | return _WrappedModel( 108 | model, self.timestep_map, self.rescale_timesteps, self.original_num_steps 109 | ) 110 | 111 | def _scale_timesteps(self, t): 112 | # Scaling is done by the wrapped model. 113 | return t 114 | 115 | 116 | class _WrappedModel: 117 | def __init__(self, model, timestep_map, rescale_timesteps, original_num_steps): 118 | self.model = model 119 | self.timestep_map = timestep_map 120 | self.rescale_timesteps = rescale_timesteps 121 | self.original_num_steps = original_num_steps 122 | 123 | def __call__(self, x, ts, **kwargs): 124 | map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype) 125 | new_ts = map_tensor[ts] 126 | if self.rescale_timesteps: 127 | new_ts = new_ts.float() * (1000.0 / self.original_num_steps) 128 | return self.model(x, new_ts, **kwargs) 129 | -------------------------------------------------------------------------------- /guided_diffusion/nn.py: -------------------------------------------------------------------------------- 1 | """ 2 | Various utilities for neural networks. 3 | """ 4 | 5 | import math 6 | 7 | import torch as th 8 | import torch.nn as nn 9 | 10 | 11 | # PyTorch 1.7 has SiLU, but we support PyTorch 1.5. 12 | class SiLU(nn.Module): 13 | def forward(self, x): 14 | return x * th.sigmoid(x) 15 | 16 | 17 | class GroupNorm32(nn.GroupNorm): 18 | def forward(self, x): 19 | return super().forward(x.float()).type(x.dtype) 20 | 21 | 22 | def conv_nd(dims, *args, **kwargs): 23 | """ 24 | Create a 1D, 2D, or 3D convolution module. 25 | """ 26 | if dims == 1: 27 | return nn.Conv1d(*args, **kwargs) 28 | elif dims == 2: 29 | return nn.Conv2d(*args, **kwargs) 30 | elif dims == 3: 31 | return nn.Conv3d(*args, **kwargs) 32 | raise ValueError(f"unsupported dimensions: {dims}") 33 | 34 | 35 | def linear(*args, **kwargs): 36 | """ 37 | Create a linear module. 38 | """ 39 | return nn.Linear(*args, **kwargs) 40 | 41 | 42 | def avg_pool_nd(dims, *args, **kwargs): 43 | """ 44 | Create a 1D, 2D, or 3D average pooling module. 45 | """ 46 | if dims == 1: 47 | return nn.AvgPool1d(*args, **kwargs) 48 | elif dims == 2: 49 | return nn.AvgPool2d(*args, **kwargs) 50 | elif dims == 3: 51 | return nn.AvgPool3d(*args, **kwargs) 52 | raise ValueError(f"unsupported dimensions: {dims}") 53 | 54 | 55 | def update_ema(target_params, source_params, rate=0.99): 56 | """ 57 | Update target parameters to be closer to those of source parameters using 58 | an exponential moving average. 59 | 60 | :param target_params: the target parameter sequence. 61 | :param source_params: the source parameter sequence. 62 | :param rate: the EMA rate (closer to 1 means slower). 63 | """ 64 | for targ, src in zip(target_params, source_params): 65 | targ.detach().mul_(rate).add_(src, alpha=1 - rate) 66 | 67 | 68 | def zero_module(module): 69 | """ 70 | Zero out the parameters of a module and return it. 71 | """ 72 | for p in module.parameters(): 73 | p.detach().zero_() 74 | return module 75 | 76 | 77 | def scale_module(module, scale): 78 | """ 79 | Scale the parameters of a module and return it. 80 | """ 81 | for p in module.parameters(): 82 | p.detach().mul_(scale) 83 | return module 84 | 85 | 86 | def mean_flat(tensor): 87 | """ 88 | Take the mean over all non-batch dimensions. 89 | """ 90 | return tensor.mean(dim=list(range(1, len(tensor.shape)))) 91 | 92 | 93 | def normalization(channels): 94 | """ 95 | Make a standard normalization layer. 96 | 97 | :param channels: number of input channels. 98 | :return: an nn.Module for normalization. 99 | """ 100 | return GroupNorm32(32, channels) 101 | 102 | 103 | def timestep_embedding(timesteps, dim, max_period=10000): 104 | """ 105 | Create sinusoidal timestep embeddings. 106 | 107 | :param timesteps: a 1-D Tensor of N indices, one per batch element. 108 | These may be fractional. 109 | :param dim: the dimension of the output. 110 | :param max_period: controls the minimum frequency of the embeddings. 111 | :return: an [N x dim] Tensor of positional embeddings. 112 | """ 113 | half = dim // 2 114 | freqs = th.exp( 115 | -math.log(max_period) * th.arange(start=0, end=half, dtype=th.float32) / half 116 | ).to(device=timesteps.device) 117 | args = timesteps[:, None].float() * freqs[None] 118 | embedding = th.cat([th.cos(args), th.sin(args)], dim=-1) 119 | if dim % 2: 120 | embedding = th.cat([embedding, th.zeros_like(embedding[:, :1])], dim=-1) 121 | return embedding 122 | 123 | 124 | def checkpoint(func, inputs, params, flag): 125 | """ 126 | Evaluate a function without caching intermediate activations, allowing for 127 | reduced memory at the expense of extra compute in the backward pass. 128 | 129 | :param func: the function to evaluate. 130 | :param inputs: the argument sequence to pass to `func`. 131 | :param params: a sequence of parameters `func` depends on but does not 132 | explicitly take as arguments. 133 | :param flag: if False, disable gradient checkpointing. 134 | """ 135 | # if flag: 136 | # args = tuple(inputs) + tuple(params) 137 | # return CheckpointFunction.apply(func, len(inputs), *args) 138 | # else: 139 | # return func(*inputs) 140 | 141 | return func(*inputs) 142 | 143 | 144 | 145 | class CheckpointFunction(th.autograd.Function): 146 | @staticmethod 147 | def forward(ctx, run_function, length, *args): 148 | ctx.run_function = run_function 149 | ctx.input_tensors = list(args[:length]) 150 | ctx.input_params = list(args[length:]) 151 | with th.no_grad(): 152 | output_tensors = ctx.run_function(*ctx.input_tensors) 153 | return output_tensors 154 | 155 | @staticmethod 156 | def backward(ctx, *output_grads): 157 | ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors] 158 | with th.enable_grad(): 159 | # Fixes a bug where the first op in run_function modifies the 160 | # Tensor storage in place, which is not allowed for detach()'d 161 | # Tensors. 162 | shallow_copies = [x.view_as(x) for x in ctx.input_tensors] 163 | output_tensors = ctx.run_function(*shallow_copies) 164 | input_grads = th.autograd.grad( 165 | output_tensors, 166 | ctx.input_tensors + ctx.input_params, 167 | output_grads, 168 | allow_unused=True, 169 | ) 170 | del ctx.input_tensors 171 | del ctx.input_params 172 | del output_tensors 173 | return (None, None) + input_grads 174 | -------------------------------------------------------------------------------- /image_sample.py: -------------------------------------------------------------------------------- 1 | """ 2 | Generate a large batch of image samples from a model and save them as a large 3 | numpy array. This can be used to produce samples for FID evaluation. 4 | """ 5 | 6 | import argparse 7 | import os 8 | 9 | import numpy as np 10 | import torch as th 11 | import torch.distributed as dist 12 | 13 | from guided_diffusion import dist_util, logger 14 | from guided_diffusion.script_util import ( 15 | NUM_CLASSES, 16 | model_and_diffusion_defaults, 17 | create_model_and_diffusion, 18 | add_dict_to_argparser, 19 | args_to_dict, 20 | ) 21 | 22 | 23 | def main(): 24 | args = create_argparser().parse_args() 25 | logger.configure(dir=args.log_dir) 26 | 27 | assert args.guidance_type in ['dps', 'dds'] 28 | my_kwargs = { 29 | 'sim_guided': args.sim_guided, 30 | 'eta': args.eta, 31 | 'exp_name': os.path.basename(args.log_dir), 32 | 'prop_dir': args.prop_dir, 33 | 'save_inter': args.save_inter, 34 | 'interval': args.interval, 35 | 'inter_rate': args.inter_rate, 36 | 'guidance_type': args.guidance_type, 37 | 'use_normed_grad': args.use_normed_grad, 38 | 'stoptime': args.stoptime, 39 | 'use_adjgrad_norm': args.use_adjgrad_norm, 40 | 'sim_type': args.sim_type, 41 | 'tsr': args.tsr, 42 | 'manual_class_id': args.manual_class_id, 43 | } 44 | print("my_kwargs: ", my_kwargs) 45 | 46 | if not args.gpu_id == '': 47 | logger.log("using device %s" % args.gpu_id) 48 | th.cuda.set_device(th.device(f"cuda:{int(args.gpu_id)}")) 49 | 50 | dist_util.setup_dist() 51 | 52 | logger.log("creating model and diffusion...") 53 | model, diffusion = create_model_and_diffusion( 54 | **args_to_dict(args, model_and_diffusion_defaults().keys()) 55 | ) 56 | model.load_state_dict( 57 | dist_util.load_state_dict(args.model_path, map_location="cpu") 58 | ) 59 | model.to(dist_util.dev()) 60 | if args.use_fp16: 61 | model.convert_to_fp16() 62 | model.eval() 63 | 64 | logger.log("sampling...") 65 | all_images = [] 66 | all_labels = [] 67 | while len(all_images) * args.batch_size < args.num_samples: 68 | model_kwargs = {} 69 | if args.class_cond: 70 | if args.manual_class_id != "": 71 | assert 0 <= int(args.manual_class_id) < args.num_classes 72 | classes = th.ones(args.batch_size, dtype=th.long, device=dist_util.dev()) * int(args.manual_class_id) 73 | else: 74 | classes = th.randint( 75 | low=0, high=args.num_classes, size=(args.batch_size,), device=dist_util.dev() 76 | ) 77 | model_kwargs["y"] = classes 78 | print("classes: ", classes) 79 | sample_fn = ( 80 | diffusion.p_sample_loop if not args.use_ddim else diffusion.ddim_sample_loop 81 | ) 82 | sample = sample_fn( 83 | model, 84 | (args.batch_size, 3 if not args.gray_imgs else 1, args.image_size, args.image_size), 85 | clip_denoised=args.clip_denoised, 86 | model_kwargs=model_kwargs, 87 | my_kwargs=my_kwargs 88 | ) 89 | sample = ((sample + 1) * 127.5).clamp(0, 255).to(th.uint8) 90 | sample = sample.permute(0, 2, 3, 1) 91 | sample = sample.contiguous() 92 | 93 | gathered_samples = [th.zeros_like(sample) for _ in range(dist.get_world_size())] 94 | dist.all_gather(gathered_samples, sample) # gather not supported with NCCL 95 | all_images.extend([sample.cpu().numpy() for sample in gathered_samples]) 96 | if args.class_cond: 97 | gathered_labels = [ 98 | th.zeros_like(classes) for _ in range(dist.get_world_size()) 99 | ] 100 | dist.all_gather(gathered_labels, classes) 101 | all_labels.extend([labels.cpu().numpy() for labels in gathered_labels]) 102 | logger.log(f"created {len(all_images) * args.batch_size} samples") 103 | 104 | arr = np.concatenate(all_images, axis=0) 105 | arr = arr[: args.num_samples] 106 | if args.class_cond: 107 | label_arr = np.concatenate(all_labels, axis=0) 108 | label_arr = label_arr[: args.num_samples] 109 | if dist.get_rank() == 0: 110 | shape_str = "x".join([str(x) for x in arr.shape]) 111 | out_path = os.path.join(logger.get_dir(), f"samples_{shape_str}.npz") 112 | logger.log(f"saving to {out_path}") 113 | if args.class_cond: 114 | np.savez(out_path, arr, label_arr) 115 | else: 116 | np.savez(out_path, arr) 117 | if args.save_img: 118 | # import imageio 119 | from PIL import Image 120 | 121 | os.makedirs(os.path.join(logger.get_dir(), "imgs"), exist_ok=True) 122 | for i, img in enumerate(arr): 123 | 124 | pil_img = Image.fromarray(img.squeeze()) 125 | pil_img.save(os.path.join(logger.get_dir(), "imgs", f"{i}.png")) 126 | # imageio.imwrite(os.path.join(logger.get_dir(), "imgs", f"{i}.png"), img) 127 | 128 | dist.barrier() 129 | logger.log("sampling complete") 130 | 131 | 132 | def create_argparser(): 133 | defaults = dict( 134 | clip_denoised=True, 135 | num_samples=10000, 136 | batch_size=16, 137 | use_ddim=False, 138 | model_path="", 139 | log_dir="./log_dir", 140 | gray_imgs=True, ##### modified by sum: for admitting grayscale images ##### 141 | num_classes=3, 142 | manual_class_id="", 143 | gpu_id="", 144 | save_img=False, 145 | sim_guided=False, # parameters for simulation-guided sampling 146 | eta=1.0, 147 | prop_dir='top', 148 | save_inter=False, 149 | interval=10, 150 | inter_rate=1, 151 | guidance_type='dps', 152 | use_normed_grad=False, 153 | stoptime=0.0, 154 | use_adjgrad_norm = False, 155 | sim_type = 'waveguide', 156 | tsr=100 157 | ) 158 | defaults.update(model_and_diffusion_defaults()) 159 | parser = argparse.ArgumentParser() 160 | add_dict_to_argparser(parser, defaults) 161 | return parser 162 | 163 | 164 | if __name__ == "__main__": 165 | main() 166 | -------------------------------------------------------------------------------- /guided_diffusion/resample.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | 3 | import numpy as np 4 | import torch as th 5 | import torch.distributed as dist 6 | 7 | 8 | def create_named_schedule_sampler(name, diffusion): 9 | """ 10 | Create a ScheduleSampler from a library of pre-defined samplers. 11 | 12 | :param name: the name of the sampler. 13 | :param diffusion: the diffusion object to sample for. 14 | """ 15 | if name == "uniform": 16 | return UniformSampler(diffusion) 17 | elif name == "loss-second-moment": 18 | return LossSecondMomentResampler(diffusion) 19 | else: 20 | raise NotImplementedError(f"unknown schedule sampler: {name}") 21 | 22 | 23 | class ScheduleSampler(ABC): 24 | """ 25 | A distribution over timesteps in the diffusion process, intended to reduce 26 | variance of the objective. 27 | 28 | By default, samplers perform unbiased importance sampling, in which the 29 | objective's mean is unchanged. 30 | However, subclasses may override sample() to change how the resampled 31 | terms are reweighted, allowing for actual changes in the objective. 32 | """ 33 | 34 | @abstractmethod 35 | def weights(self): 36 | """ 37 | Get a numpy array of weights, one per diffusion step. 38 | 39 | The weights needn't be normalized, but must be positive. 40 | """ 41 | 42 | def sample(self, batch_size, device): 43 | """ 44 | Importance-sample timesteps for a batch. 45 | 46 | :param batch_size: the number of timesteps. 47 | :param device: the torch device to save to. 48 | :return: a tuple (timesteps, weights): 49 | - timesteps: a tensor of timestep indices. 50 | - weights: a tensor of weights to scale the resulting losses. 51 | """ 52 | w = self.weights() 53 | p = w / np.sum(w) 54 | indices_np = np.random.choice(len(p), size=(batch_size,), p=p) 55 | indices = th.from_numpy(indices_np).long().to(device) 56 | weights_np = 1 / (len(p) * p[indices_np]) 57 | weights = th.from_numpy(weights_np).float().to(device) 58 | return indices, weights 59 | 60 | 61 | class UniformSampler(ScheduleSampler): 62 | def __init__(self, diffusion): 63 | self.diffusion = diffusion 64 | self._weights = np.ones([diffusion.num_timesteps]) 65 | 66 | def weights(self): 67 | return self._weights 68 | 69 | 70 | class LossAwareSampler(ScheduleSampler): 71 | def update_with_local_losses(self, local_ts, local_losses): 72 | """ 73 | Update the reweighting using losses from a model. 74 | 75 | Call this method from each rank with a batch of timesteps and the 76 | corresponding losses for each of those timesteps. 77 | This method will perform synchronization to make sure all of the ranks 78 | maintain the exact same reweighting. 79 | 80 | :param local_ts: an integer Tensor of timesteps. 81 | :param local_losses: a 1D Tensor of losses. 82 | """ 83 | batch_sizes = [ 84 | th.tensor([0], dtype=th.int32, device=local_ts.device) 85 | for _ in range(dist.get_world_size()) 86 | ] 87 | dist.all_gather( 88 | batch_sizes, 89 | th.tensor([len(local_ts)], dtype=th.int32, device=local_ts.device), 90 | ) 91 | 92 | # Pad all_gather batches to be the maximum batch size. 93 | batch_sizes = [x.item() for x in batch_sizes] 94 | max_bs = max(batch_sizes) 95 | 96 | timestep_batches = [th.zeros(max_bs).to(local_ts) for bs in batch_sizes] 97 | loss_batches = [th.zeros(max_bs).to(local_losses) for bs in batch_sizes] 98 | dist.all_gather(timestep_batches, local_ts) 99 | dist.all_gather(loss_batches, local_losses) 100 | timesteps = [ 101 | x.item() for y, bs in zip(timestep_batches, batch_sizes) for x in y[:bs] 102 | ] 103 | losses = [x.item() for y, bs in zip(loss_batches, batch_sizes) for x in y[:bs]] 104 | self.update_with_all_losses(timesteps, losses) 105 | 106 | @abstractmethod 107 | def update_with_all_losses(self, ts, losses): 108 | """ 109 | Update the reweighting using losses from a model. 110 | 111 | Sub-classes should override this method to update the reweighting 112 | using losses from the model. 113 | 114 | This method directly updates the reweighting without synchronizing 115 | between workers. It is called by update_with_local_losses from all 116 | ranks with identical arguments. Thus, it should have deterministic 117 | behavior to maintain state across workers. 118 | 119 | :param ts: a list of int timesteps. 120 | :param losses: a list of float losses, one per timestep. 121 | """ 122 | 123 | 124 | class LossSecondMomentResampler(LossAwareSampler): 125 | def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001): 126 | self.diffusion = diffusion 127 | self.history_per_term = history_per_term 128 | self.uniform_prob = uniform_prob 129 | self._loss_history = np.zeros( 130 | [diffusion.num_timesteps, history_per_term], dtype=np.float64 131 | ) 132 | self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int) 133 | 134 | def weights(self): 135 | if not self._warmed_up(): 136 | return np.ones([self.diffusion.num_timesteps], dtype=np.float64) 137 | weights = np.sqrt(np.mean(self._loss_history ** 2, axis=-1)) 138 | weights /= np.sum(weights) 139 | weights *= 1 - self.uniform_prob 140 | weights += self.uniform_prob / len(weights) 141 | return weights 142 | 143 | def update_with_all_losses(self, ts, losses): 144 | for t, loss in zip(ts, losses): 145 | if self._loss_counts[t] == self.history_per_term: 146 | # Shift out the oldest loss term. 147 | self._loss_history[t, :-1] = self._loss_history[t, 1:] 148 | self._loss_history[t, -1] = loss 149 | else: 150 | self._loss_history[t, self._loss_counts[t]] = loss 151 | self._loss_counts[t] += 1 152 | 153 | def _warmed_up(self): 154 | return (self._loss_counts == self.history_per_term).all() 155 | -------------------------------------------------------------------------------- /guided_diffusion/image_datasets.py: -------------------------------------------------------------------------------- 1 | import math 2 | import random 3 | 4 | from PIL import Image 5 | import blobfile as bf 6 | from mpi4py import MPI 7 | import numpy as np 8 | from torch.utils.data import DataLoader, Dataset 9 | 10 | 11 | def load_data( 12 | *, 13 | data_dir, 14 | batch_size, 15 | image_size, 16 | class_cond=False, 17 | deterministic=False, 18 | random_crop=False, 19 | random_flip=True, 20 | gray_imgs=False ##### modified by sum: for admitting grayscale images ##### 21 | ): 22 | """ 23 | For a dataset, create a generator over (images, kwargs) pairs. 24 | 25 | Each images is an NCHW float tensor, and the kwargs dict contains zero or 26 | more keys, each of which map to a batched Tensor of their own. 27 | The kwargs dict can be used for class labels, in which case the key is "y" 28 | and the values are integer tensors of class labels. 29 | 30 | :param data_dir: a dataset directory. 31 | :param batch_size: the batch size of each returned pair. 32 | :param image_size: the size to which images are resized. 33 | :param class_cond: if True, include a "y" key in returned dicts for class 34 | label. If classes are not available and this is true, an 35 | exception will be raised. 36 | :param deterministic: if True, yield results in a deterministic order. 37 | :param random_crop: if True, randomly crop the images for augmentation. 38 | :param random_flip: if True, randomly flip the images for augmentation. 39 | """ 40 | if not data_dir: 41 | raise ValueError("unspecified data directory") 42 | all_files = _list_image_files_recursively(data_dir) 43 | classes = None 44 | if class_cond: 45 | # Assume classes are the first part of the filename, 46 | # before an underscore. 47 | class_names = [bf.basename(path).split("_")[0] for path in all_files] 48 | sorted_classes = {x: i for i, x in enumerate(sorted(set(class_names)))} 49 | classes = [sorted_classes[x] for x in class_names] 50 | dataset = ImageDataset( 51 | image_size, 52 | all_files, 53 | classes=classes, 54 | shard=MPI.COMM_WORLD.Get_rank(), 55 | num_shards=MPI.COMM_WORLD.Get_size(), 56 | random_crop=random_crop, 57 | random_flip=random_flip, 58 | gray_imgs=gray_imgs, ##### modified by sum: for admitting grayscale images ##### 59 | ) 60 | if deterministic: 61 | loader = DataLoader( 62 | dataset, batch_size=batch_size, shuffle=False, num_workers=1, drop_last=True 63 | ) 64 | else: 65 | loader = DataLoader( 66 | dataset, batch_size=batch_size, shuffle=True, num_workers=1, drop_last=True 67 | ) 68 | while True: 69 | yield from loader 70 | 71 | 72 | def _list_image_files_recursively(data_dir): 73 | results = [] 74 | for entry in sorted(bf.listdir(data_dir)): 75 | full_path = bf.join(data_dir, entry) 76 | ext = entry.split(".")[-1] 77 | if "." in entry and ext.lower() in ["jpg", "jpeg", "png", "gif"]: 78 | results.append(full_path) 79 | elif bf.isdir(full_path): 80 | results.extend(_list_image_files_recursively(full_path)) 81 | return results 82 | 83 | 84 | class ImageDataset(Dataset): 85 | def __init__( 86 | self, 87 | resolution, 88 | image_paths, 89 | classes=None, 90 | shard=0, 91 | num_shards=1, 92 | random_crop=False, 93 | random_flip=True, 94 | gray_imgs=False, ##### modified by sum: for admitting grayscale images ##### 95 | ): 96 | super().__init__() 97 | self.resolution = resolution 98 | self.local_images = image_paths[shard:][::num_shards] 99 | self.local_classes = None if classes is None else classes[shard:][::num_shards] 100 | self.random_crop = random_crop 101 | self.random_flip = random_flip 102 | self.gray_imgs = gray_imgs 103 | 104 | def __len__(self): 105 | return len(self.local_images) 106 | 107 | def __getitem__(self, idx): 108 | path = self.local_images[idx] 109 | with bf.BlobFile(path, "rb") as f: 110 | pil_image = Image.open(f) 111 | pil_image.load() 112 | if not self.gray_imgs: ##### modified by sum: for admitting grayscale images ##### 113 | pil_image = pil_image.convert("RGB") 114 | else: ##### modified by sum: for admitting grayscale images ##### 115 | pil_image = pil_image.convert("L") 116 | 117 | if self.random_crop: 118 | arr = random_crop_arr(pil_image, self.resolution, self.gray_imgs) 119 | else: 120 | arr = center_crop_arr(pil_image, self.resolution, self.gray_imgs) 121 | 122 | 123 | if self.random_flip and random.random() < 0.5: 124 | arr = arr[:, ::-1] 125 | 126 | arr = arr.astype(np.float32) / 127.5 - 1 127 | 128 | out_dict = {} 129 | if self.local_classes is not None: 130 | out_dict["y"] = np.array(self.local_classes[idx], dtype=np.int64) 131 | return np.transpose(arr, [2, 0, 1]), out_dict 132 | 133 | 134 | def center_crop_arr(pil_image, image_size, gray_imgs=False): 135 | # We are not on a new enough PIL to support the `reducing_gap` 136 | # argument, which uses BOX downsampling at powers of two first. 137 | # Thus, we do it by hand to improve downsample quality. 138 | while min(*pil_image.size) >= 2 * image_size: 139 | pil_image = pil_image.resize( 140 | tuple(x // 2 for x in pil_image.size), resample=Image.BOX 141 | ) 142 | 143 | scale = image_size / min(*pil_image.size) 144 | pil_image = pil_image.resize( 145 | tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC 146 | ) 147 | 148 | arr = np.array(pil_image) 149 | if gray_imgs: ##### modified by sum: for admitting grayscale images ##### 150 | arr = arr[..., None] 151 | crop_y = (arr.shape[0] - image_size) // 2 152 | crop_x = (arr.shape[1] - image_size) // 2 153 | return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size] 154 | 155 | 156 | def random_crop_arr(pil_image, image_size, min_crop_frac=0.8, max_crop_frac=1.0, gray_imgs=False): 157 | min_smaller_dim_size = math.ceil(image_size / max_crop_frac) 158 | max_smaller_dim_size = math.ceil(image_size / min_crop_frac) 159 | smaller_dim_size = random.randrange(min_smaller_dim_size, max_smaller_dim_size + 1) 160 | 161 | # We are not on a new enough PIL to support the `reducing_gap` 162 | # argument, which uses BOX downsampling at powers of two first. 163 | # Thus, we do it by hand to improve downsample quality. 164 | while min(*pil_image.size) >= 2 * smaller_dim_size: 165 | pil_image = pil_image.resize( 166 | tuple(x // 2 for x in pil_image.size), resample=Image.BOX 167 | ) 168 | 169 | scale = smaller_dim_size / min(*pil_image.size) 170 | pil_image = pil_image.resize( 171 | tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC 172 | ) 173 | 174 | arr = np.array(pil_image) 175 | if gray_imgs: ##### modified by sum: for admitting grayscale images ##### 176 | arr = arr[..., None] 177 | crop_y = random.randrange(arr.shape[0] - image_size + 1) 178 | crop_x = random.randrange(arr.shape[1] - image_size + 1) 179 | return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size] 180 | -------------------------------------------------------------------------------- /guided_diffusion/fp16_util.py: -------------------------------------------------------------------------------- 1 | """ 2 | Helpers to train with 16-bit precision. 3 | """ 4 | 5 | import numpy as np 6 | import torch as th 7 | import torch.nn as nn 8 | from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors 9 | 10 | from . import logger 11 | 12 | INITIAL_LOG_LOSS_SCALE = 20.0 13 | 14 | 15 | def convert_module_to_f16(l): 16 | """ 17 | Convert primitive modules to float16. 18 | """ 19 | if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)): 20 | l.weight.data = l.weight.data.half() 21 | if l.bias is not None: 22 | l.bias.data = l.bias.data.half() 23 | 24 | 25 | def convert_module_to_f32(l): 26 | """ 27 | Convert primitive modules to float32, undoing convert_module_to_f16(). 28 | """ 29 | if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)): 30 | l.weight.data = l.weight.data.float() 31 | if l.bias is not None: 32 | l.bias.data = l.bias.data.float() 33 | 34 | 35 | def make_master_params(param_groups_and_shapes): 36 | """ 37 | Copy model parameters into a (differently-shaped) list of full-precision 38 | parameters. 39 | """ 40 | master_params = [] 41 | for param_group, shape in param_groups_and_shapes: 42 | master_param = nn.Parameter( 43 | _flatten_dense_tensors( 44 | [param.detach().float() for (_, param) in param_group] 45 | ).view(shape) 46 | ) 47 | master_param.requires_grad = True 48 | master_params.append(master_param) 49 | return master_params 50 | 51 | 52 | def model_grads_to_master_grads(param_groups_and_shapes, master_params): 53 | """ 54 | Copy the gradients from the model parameters into the master parameters 55 | from make_master_params(). 56 | """ 57 | for master_param, (param_group, shape) in zip( 58 | master_params, param_groups_and_shapes 59 | ): 60 | master_param.grad = _flatten_dense_tensors( 61 | [param_grad_or_zeros(param) for (_, param) in param_group] 62 | ).view(shape) 63 | 64 | 65 | def master_params_to_model_params(param_groups_and_shapes, master_params): 66 | """ 67 | Copy the master parameter data back into the model parameters. 68 | """ 69 | # Without copying to a list, if a generator is passed, this will 70 | # silently not copy any parameters. 71 | for master_param, (param_group, _) in zip(master_params, param_groups_and_shapes): 72 | for (_, param), unflat_master_param in zip( 73 | param_group, unflatten_master_params(param_group, master_param.view(-1)) 74 | ): 75 | param.detach().copy_(unflat_master_param) 76 | 77 | 78 | def unflatten_master_params(param_group, master_param): 79 | return _unflatten_dense_tensors(master_param, [param for (_, param) in param_group]) 80 | 81 | 82 | def get_param_groups_and_shapes(named_model_params): 83 | named_model_params = list(named_model_params) 84 | scalar_vector_named_params = ( 85 | [(n, p) for (n, p) in named_model_params if p.ndim <= 1], 86 | (-1), 87 | ) 88 | matrix_named_params = ( 89 | [(n, p) for (n, p) in named_model_params if p.ndim > 1], 90 | (1, -1), 91 | ) 92 | return [scalar_vector_named_params, matrix_named_params] 93 | 94 | 95 | def master_params_to_state_dict( 96 | model, param_groups_and_shapes, master_params, use_fp16 97 | ): 98 | if use_fp16: 99 | state_dict = model.state_dict() 100 | for master_param, (param_group, _) in zip( 101 | master_params, param_groups_and_shapes 102 | ): 103 | for (name, _), unflat_master_param in zip( 104 | param_group, unflatten_master_params(param_group, master_param.view(-1)) 105 | ): 106 | assert name in state_dict 107 | state_dict[name] = unflat_master_param 108 | else: 109 | state_dict = model.state_dict() 110 | for i, (name, _value) in enumerate(model.named_parameters()): 111 | assert name in state_dict 112 | state_dict[name] = master_params[i] 113 | return state_dict 114 | 115 | 116 | def state_dict_to_master_params(model, state_dict, use_fp16): 117 | if use_fp16: 118 | named_model_params = [ 119 | (name, state_dict[name]) for name, _ in model.named_parameters() 120 | ] 121 | param_groups_and_shapes = get_param_groups_and_shapes(named_model_params) 122 | master_params = make_master_params(param_groups_and_shapes) 123 | else: 124 | master_params = [state_dict[name] for name, _ in model.named_parameters()] 125 | return master_params 126 | 127 | 128 | def zero_master_grads(master_params): 129 | for param in master_params: 130 | param.grad = None 131 | 132 | 133 | def zero_grad(model_params): 134 | for param in model_params: 135 | # Taken from https://pytorch.org/docs/stable/_modules/torch/optim/optimizer.html#Optimizer.add_param_group 136 | if param.grad is not None: 137 | param.grad.detach_() 138 | param.grad.zero_() 139 | 140 | 141 | def param_grad_or_zeros(param): 142 | if param.grad is not None: 143 | return param.grad.data.detach() 144 | else: 145 | return th.zeros_like(param) 146 | 147 | 148 | class MixedPrecisionTrainer: 149 | def __init__( 150 | self, 151 | *, 152 | model, 153 | use_fp16=False, 154 | fp16_scale_growth=1e-3, 155 | initial_lg_loss_scale=INITIAL_LOG_LOSS_SCALE, 156 | ): 157 | self.model = model 158 | self.use_fp16 = use_fp16 159 | self.fp16_scale_growth = fp16_scale_growth 160 | 161 | self.model_params = list(self.model.parameters()) 162 | self.master_params = self.model_params 163 | self.param_groups_and_shapes = None 164 | self.lg_loss_scale = initial_lg_loss_scale 165 | 166 | if self.use_fp16: 167 | self.param_groups_and_shapes = get_param_groups_and_shapes( 168 | self.model.named_parameters() 169 | ) 170 | self.master_params = make_master_params(self.param_groups_and_shapes) 171 | self.model.convert_to_fp16() 172 | 173 | def zero_grad(self): 174 | zero_grad(self.model_params) 175 | 176 | def backward(self, loss: th.Tensor): 177 | if self.use_fp16: 178 | loss_scale = 2 ** self.lg_loss_scale 179 | (loss * loss_scale).backward() 180 | else: 181 | loss.backward() 182 | 183 | def optimize(self, opt: th.optim.Optimizer): 184 | if self.use_fp16: 185 | return self._optimize_fp16(opt) 186 | else: 187 | return self._optimize_normal(opt) 188 | 189 | def _optimize_fp16(self, opt: th.optim.Optimizer): 190 | logger.logkv_mean("lg_loss_scale", self.lg_loss_scale) 191 | model_grads_to_master_grads(self.param_groups_and_shapes, self.master_params) 192 | grad_norm, param_norm = self._compute_norms(grad_scale=2 ** self.lg_loss_scale) 193 | if check_overflow(grad_norm): 194 | self.lg_loss_scale -= 1 195 | logger.log(f"Found NaN, decreased lg_loss_scale to {self.lg_loss_scale}") 196 | zero_master_grads(self.master_params) 197 | return False 198 | 199 | logger.logkv_mean("grad_norm", grad_norm) 200 | logger.logkv_mean("param_norm", param_norm) 201 | 202 | for p in self.master_params: 203 | p.grad.mul_(1.0 / (2 ** self.lg_loss_scale)) 204 | opt.step() 205 | zero_master_grads(self.master_params) 206 | master_params_to_model_params(self.param_groups_and_shapes, self.master_params) 207 | self.lg_loss_scale += self.fp16_scale_growth 208 | return True 209 | 210 | def _optimize_normal(self, opt: th.optim.Optimizer): 211 | grad_norm, param_norm = self._compute_norms() 212 | logger.logkv_mean("grad_norm", grad_norm) 213 | logger.logkv_mean("param_norm", param_norm) 214 | opt.step() 215 | return True 216 | 217 | def _compute_norms(self, grad_scale=1.0): 218 | grad_norm = 0.0 219 | param_norm = 0.0 220 | for p in self.master_params: 221 | with th.no_grad(): 222 | param_norm += th.norm(p, p=2, dtype=th.float32).item() ** 2 223 | if p.grad is not None: 224 | grad_norm += th.norm(p.grad, p=2, dtype=th.float32).item() ** 2 225 | return np.sqrt(grad_norm) / grad_scale, np.sqrt(param_norm) 226 | 227 | def master_params_to_state_dict(self, master_params): 228 | return master_params_to_state_dict( 229 | self.model, self.param_groups_and_shapes, master_params, self.use_fp16 230 | ) 231 | 232 | def state_dict_to_master_params(self, state_dict): 233 | return state_dict_to_master_params(self.model, state_dict, self.use_fp16) 234 | 235 | 236 | def check_overflow(value): 237 | return (value == float("inf")) or (value == -float("inf")) or (value != value) 238 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # AdjointDiffusion 2 | 3 | [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE) [![Python: 3.9+](https://img.shields.io/badge/python-3.9%2B-blue.svg)](https://www.python.org/) 4 | 5 | 6 | ![logo](images/logo.png) 7 | 8 | **AdjointDiffusion** is a new method for structural optimization using diffusion models. 9 | It is a **physics-guided and fabrication-aware structural optimization** using diffusion models augmented with adjoint gradient. By combining powerful generative models with adjoint sensitivity analysis, this approach can more efficiently discover complex, high-performance designs than the traditional methods. 10 | 11 | The codes are provided following the paper named [Physics-guided and fabrication-aware 12 | structural optimization using diffusion models](https://arxiv.org/abs/2504.17077) 13 | 14 | --- 15 | 16 | ## Table of Contents 17 | 18 | 1. [TL;DR](#TL;DR) 19 | 2. [Intuitive Explanation of Diffusion Models](#intuitive-explanation-of-diffusion-models) 20 | 3. [Installation](#installation) 21 | 4. [Quick Start](#quick-start) 22 | 5. [Customize Your Simulation](#customize-your-simulation) 23 | 6. [Experiment Logging with Weights & Biases](#experiment-logging-with-weights--biases) 24 | 7. [Results](#results) 25 | 8. [Code Organization](#code-organization) 26 | 9. [Citation](#citation) 27 | 28 | ## TL;DR 29 | 30 | ✨ **Integrating adjoint sensitivity analysis with diffusion models can generate high-performance and interesting structures!** 31 | 32 | ![comparison](images/comparison.png) 33 | 34 | 35 | Key features: 36 | - **Adjoint Sensitivity Integration**: Seamlessly incorporates adjoint gradients into the diffusion process. 37 | - **Fabrication Constraints**: Accounts for manufacturability, ensuring real-world feasibility. 38 | - **Extensibility**: Users can use their own datasets or simulations. 39 | - **Experiment Tracking & Visualization**: Integrates with [Weights & Biases](https://wandb.ai/home). 40 | 41 | --- 42 | 43 | ## Intuitive Explanation of Diffusion Models 44 | 45 | Imagine an ink drop falling into water — it slowly spreads and dissolves. Diffusion models mimic this process in reverse: they start from noise and slowly form meaningful structures. By guiding this "reverse diffusion" with gradients from an adjoint method, we ensure the final designs are optimized and fabrication-ready. 46 | 47 | --- 48 | 49 | 50 | 51 | ## Installation 52 | 53 | This setup ensures compatibility between Meep and PyTorch. If you find any alternatives, feel free to contribute improvements via pull requests! 54 | 55 | ### 1. Clone the repository 56 | ```bash 57 | git clone https://github.com/dongjin-seo2020/AdjointDiffusion.git 58 | cd AdjointDiffusion 59 | ``` 60 | 61 | ### 2. Set up a Python environment (recommended) 62 | 63 | 64 | #### Using conda 65 | To create and activate the recommended environment with necessary dependencies: 66 | 67 | ```bash 68 | conda create -n adjoint_diffusion -c conda-forge pymeep pymeep-extras python=3.9 69 | conda activate adjoint_diffusion 70 | ``` 71 | 72 | 73 | ### 3. Installation of other libraries 74 | 75 | Install `torch` following the command (recommended): 76 | ```bash 77 | pip install torch==2.0.1+cu117 torchvision==0.15.2+cu117 --index-url https://download.pytorch.org/whl/cu117 78 | ``` 79 | 80 | Note: It has been observed that newer NVIDIA GPUs (e.g., RTX 5090) may not be compatible with this specific version of PyTorch. If you encounter issues, please refer to the official PyTorch installation guide to find a version compatible with your hardware: 81 | https://pytorch.org/get-started/locally/ 82 | 83 | Install the required packages listed in `requirements.txt`: 84 | ```bash 85 | pip install -r requirements.txt 86 | ``` 87 | 88 | 89 | ### 4. Error Handling 90 | 91 | 92 | If you encounter permission-related issues when trying to run the training script, make sure it is executable by running: 93 | 94 | ```bash 95 | chmod +x 01-train.sh 96 | ``` 97 | 98 | The, you can execute it with: 99 | 100 | ```bash 101 | ./01-train.sh 102 | ``` 103 | 104 | If you encounter errors while installing `mpi4py`, try the following steps: 105 | ```bash 106 | apt --fix-broken install 107 | apt install mpich 108 | pip install mpi4py 109 | ``` 110 | Make sure you have root access when using `apt`. 111 | 112 | 113 | 114 | --- 115 | 116 | ## Quick Start 117 | 118 | 1. **Generate a dataset:** 119 | 120 | ```bash 121 | python dataset_generation.py 122 | ``` 123 | 124 | - The data will be saved at `datasets//sigma/struct/`, where `n` is the structure dimension (e.g., `n=64` generates 64×64 binary structures) and `k` is the variance of the Gaussian filter (a larger `k` increases the minimum feature size). 125 | - Note: To reproduce the condition in the paper, run the code for `k=2`, `k=5`, and `k=8` (three times). Or you can download our pretrained network from https://zenodo.org/records/15399997 126 | - Note: You can also use your own dataset here! Provide your fabrication-satisfying image dataset and train the diffusion model with it! 127 | 128 | 2. **Update the training and sampling scripts** to specify the appropriate output directories. 129 | 130 | - For example, for `train.sh`, you should specify the variables as: 131 | ```sh 132 | DATA_DIR=/path/to/datasets 133 | LOG_DIR=path/to/experiments 134 | GPU_ID=0 135 | ``` 136 | 137 | - Or, you can set environment variables (Linux/macOS): 138 | ```bash 139 | export DATA_DIR=/path/to/datasets 140 | export LOG_DIR=/path/to/experiments 141 | export GPU_ID=0 142 | ``` 143 | 144 | - For detailed usage examples, including training and sampling with actual settings, see: 145 | - [training_examples.ipynb](./training_examples.ipynb) 146 | - [sampling_example.ipynb](./sampling_example.ipynb) 147 | 148 | 3. **Train a diffusion model:** 149 | 150 | ```bash 151 | ./01-train.sh 152 | ``` 153 | 154 | - Alternatively: run `02-train.ipynb` 155 | - Note: Set `--class_cond` to `False` if your dataset contains only a single structural condition (i.e., no class conditioning needed). 156 | If you have multiple structural conditions (e.g., different fabrication constraints), set it to `True` to enable class-conditional training. 157 | - Note: The training process will continue indefinitely unless manually stopped. In our setup, training for around 25,000 steps produced satisfactory results, though fewer steps may also be sufficient. If you are using a customized dataset, the optimal number of steps may vary. 158 | 159 | 160 | 4. **Sample and optimize structures:** 161 | 162 | ```bash 163 | ./01-sample.sh 164 | ``` 165 | 166 | - Alternatively: run `02-sample.ipynb` 167 | - Note: Set `--class_cond` to `False` if your dataset contains only a single structural condition (i.e., no class conditioning needed). 168 | If you have multiple structural conditions (e.g., different fabrication constraints), set it to `True` to enable class-conditional training. 169 | - Note: We recommend to use the model where the name is `ema_0.9999_*.pt`. The number in `*` means the training step number. `ema` is for `Exponential Moving Average`. 170 | 171 | 172 | 5. **View outputs** 173 | - Every output (performance, structure) is logged in [wandb](#experiment-logging-with-weights--biases). 174 | - Logs and generated structures are saved in `./logs/` 175 | 176 | 177 | 6. **Baseline Algorithms** 178 | 179 | We provide baseline algorithms in the `./baseline_algorithms` directory. These include **nlopt** methods like MMA for comparison. 180 | 181 | 182 | --- 183 | 184 | ## Customize Your Simulation 185 | 186 | If you'd like to **integrate a custom physical simulation** into the reverse diffusion process, follow these steps: 187 | 188 | 1. **Implement Your simulation class** in `guided_diffusion/simulation.py`. 189 | 190 | Create a class that defines how to compute the figure of merit (FoM) and its corresponding adjoint gradient. For example: 191 | ```python 192 | class YourSimClass: 193 | def __init__(self, ...): 194 | ... 195 | def compute_fom(self, structure): 196 | ... 197 | def compute_adjoint(self, structure): 198 | ... 199 | ``` 200 | 201 | 202 | 2. **Update the import** in `guided_diffusion/gaussian_diffusion.py`: 203 | 204 | Replace the existing simulation import with your custom class: 205 | ```python 206 | from guided_diffusion.simulation import YourSimClass 207 | ``` 208 | 209 | 3. **Plug your simulation into the sampling loop** in `guided_diffusion/gaussian_diffusion.py` 210 | In `guided_diffusion/gaussian_diffusion.py`, locate where `simulation_()` is called (typically inside the `p_sample()` function), and replace it with your custom simulation logic. Make sure your class is initialized properly and passed via `my_kwargs`. 211 | For example, for `my_kwargs`: 212 | ```python 213 | my_kwargs = { 214 | "sim_guided": True, 215 | "simulation_": YourSimClass(...), 216 | "eta": 0.05, 217 | "inter_rate": 25, 218 | "stoptime": 0.1, 219 | "guidance_type": "dps", # or "dds" 220 | "exp_name": "experiment1", 221 | ... 222 | } 223 | ``` 224 | 225 | Now, your custom simulation will be used during the reverse diffusion process. 226 | 227 | --- 228 | 229 | 230 | ## Experiment Logging with Weights & Biases 231 | 232 | We use [wandb](https://wandb.ai/home) for logging and visualization. 233 | 234 | 1. Sign up at [wandb.ai](https://wandb.ai) 235 | 2. Log in: 236 | ```bash 237 | wandb login 238 | ``` 239 | 3. Run any training/sampling script and it will automatically log data to wandb. 240 | 241 | --- 242 | 243 | 244 | 245 | ## Results 246 | 247 | We visualize the performance of AdjointDiffusion across different tasks and configurations. 248 | 249 | 250 | ### Optimization Convergence and Comparisons - Problem Setup 1 (Waveguide) 251 | 252 | ![Performance Plot 1](images/Result1.png) 253 | 254 | 255 | 256 | 257 | ### Comparison of Generated Structures - Problem Setup 1 (Waveguide) 258 | 259 | ![Bar Plots](images/Result1-2.png) 260 | 261 | 262 | ### Optimization Convergence and Comparisons 2 - Problem Setup 2 (Multi-wavelength Splitter) 263 | 264 | ![Performance Plot 2](images/Result2.png) 265 | 266 | ### Comparison of Generated Structures - Problem Setup 2 (Multi-wavelength Splitter) 267 | 268 | ![Bar Plots](images/result2-2.png) 269 | 270 | 271 | ### Result: Waveguide 272 | 273 | ![Waveguide](images/Result-waveguide.png) 274 | 275 | 276 | ### Result: Multi-wavelength Splitter 277 | 278 | ![Color Router](images/Result-colorrouter.png) 279 | 280 | --- 281 | 282 | ## Code Organization 283 | 284 | ``` 285 | AdjointDiffusion/ 286 | ├── dataset_generation.py # Dataset generation script 287 | ├── image_train.py # Main training script 288 | ├── image_sample.py # Main sampling script 289 | ├── requirements.txt # Python dependencies 290 | ├── guided_diffusion/ # Backend of diffusion models 291 | └── baseline_algorithms/ # Baseline algorithms (nlopt, Gradient Ascent) 292 | ``` 293 | 294 | --- 295 | 296 | ## Citation 297 | 298 | If you use this code, please cite the following paper: 299 | 300 | ```bibtex 301 | @article{YourCitation, 302 | title = {Physics-guided and fabrication-aware structural optimization using diffusion models}, 303 | author = {Dongjin Seo†, Soobin Um†, Sangbin Lee, Jong Chul Ye*, Haejun Chung*}, 304 | journal = {arXiv}, 305 | year = {2025}, 306 | url = {https://arxiv.org/abs/2504.17077} 307 | } 308 | ``` 309 | 310 | --- 311 | 312 | 313 | ## Acknowledgements 314 | 315 | Parts of this repository are adapted from [OpenAI's guided-diffusion](https://github.com/openai/guided-diffusion), which is licensed under the MIT License. 316 | 317 | We thank the OpenAI team for their contribution. Significant modifications have been made to enable adjoint sensitivity integration and fabrication-aware optimization. 318 | 319 | --- 320 | **Happy Diffusing & Optimizing!** 321 | -------------------------------------------------------------------------------- /guided_diffusion/train_util.py: -------------------------------------------------------------------------------- 1 | import copy 2 | import functools 3 | import os 4 | 5 | import blobfile as bf 6 | import torch as th 7 | import torch.distributed as dist 8 | from torch.nn.parallel.distributed import DistributedDataParallel as DDP 9 | from torch.optim import AdamW 10 | 11 | from . import dist_util, logger 12 | from .fp16_util import MixedPrecisionTrainer 13 | from .nn import update_ema 14 | from .resample import LossAwareSampler, UniformSampler 15 | 16 | import time 17 | 18 | # For ImageNet experiments, this was a good default value. 19 | # We found that the lg_loss_scale quickly climbed to 20 | # 20-21 within the first ~1K steps of training. 21 | INITIAL_LOG_LOSS_SCALE = 20.0 22 | 23 | 24 | class TrainLoop: 25 | def __init__( 26 | self, 27 | *, 28 | model, 29 | diffusion, 30 | data, 31 | batch_size, 32 | microbatch, 33 | lr, 34 | ema_rate, 35 | log_interval, 36 | save_interval, 37 | resume_checkpoint, 38 | use_fp16=False, 39 | fp16_scale_growth=1e-3, 40 | schedule_sampler=None, 41 | weight_decay=0.0, 42 | lr_anneal_steps=0, 43 | ): 44 | self.model = model 45 | self.diffusion = diffusion 46 | self.data = data 47 | self.batch_size = batch_size 48 | self.microbatch = microbatch if microbatch > 0 else batch_size 49 | self.lr = lr 50 | self.ema_rate = ( 51 | [ema_rate] 52 | if isinstance(ema_rate, float) 53 | else [float(x) for x in ema_rate.split(",")] 54 | ) 55 | self.log_interval = log_interval 56 | self.save_interval = save_interval 57 | self.resume_checkpoint = resume_checkpoint 58 | self.use_fp16 = use_fp16 59 | self.fp16_scale_growth = fp16_scale_growth 60 | self.schedule_sampler = schedule_sampler or UniformSampler(diffusion) 61 | self.weight_decay = weight_decay 62 | self.lr_anneal_steps = lr_anneal_steps 63 | 64 | self.step = 0 65 | self.resume_step = 0 66 | self.global_batch = self.batch_size * dist.get_world_size() 67 | 68 | self.sync_cuda = th.cuda.is_available() 69 | 70 | self._load_and_sync_parameters() 71 | self.mp_trainer = MixedPrecisionTrainer( 72 | model=self.model, 73 | use_fp16=self.use_fp16, 74 | fp16_scale_growth=fp16_scale_growth, 75 | ) 76 | 77 | self.opt = AdamW( 78 | self.mp_trainer.master_params, lr=self.lr, weight_decay=self.weight_decay 79 | ) 80 | if self.resume_step: 81 | self._load_optimizer_state() 82 | # Model was resumed, either due to a restart or a checkpoint 83 | # being specified at the command line. 84 | self.ema_params = [ 85 | self._load_ema_parameters(rate) for rate in self.ema_rate 86 | ] 87 | else: 88 | self.ema_params = [ 89 | copy.deepcopy(self.mp_trainer.master_params) 90 | for _ in range(len(self.ema_rate)) 91 | ] 92 | 93 | if th.cuda.is_available(): 94 | self.use_ddp = True 95 | self.ddp_model = DDP( 96 | self.model, 97 | device_ids=[dist_util.dev()], 98 | output_device=dist_util.dev(), 99 | broadcast_buffers=False, 100 | bucket_cap_mb=128, 101 | find_unused_parameters=False, 102 | ) 103 | else: 104 | if dist.get_world_size() > 1: 105 | logger.warn( 106 | "Distributed training requires CUDA. " 107 | "Gradients will not be synchronized properly!" 108 | ) 109 | self.use_ddp = False 110 | self.ddp_model = self.model 111 | 112 | def _load_and_sync_parameters(self): 113 | resume_checkpoint = find_resume_checkpoint() or self.resume_checkpoint 114 | 115 | if resume_checkpoint: 116 | self.resume_step = parse_resume_step_from_filename(resume_checkpoint) 117 | if dist.get_rank() == 0: 118 | logger.log(f"loading model from checkpoint: {resume_checkpoint}...") 119 | self.model.load_state_dict( 120 | dist_util.load_state_dict( 121 | resume_checkpoint, map_location=dist_util.dev() 122 | ) 123 | ) 124 | 125 | dist_util.sync_params(self.model.parameters()) 126 | 127 | def _load_ema_parameters(self, rate): 128 | ema_params = copy.deepcopy(self.mp_trainer.master_params) 129 | 130 | main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint 131 | ema_checkpoint = find_ema_checkpoint(main_checkpoint, self.resume_step, rate) 132 | if ema_checkpoint: 133 | if dist.get_rank() == 0: 134 | logger.log(f"loading EMA from checkpoint: {ema_checkpoint}...") 135 | state_dict = dist_util.load_state_dict( 136 | ema_checkpoint, map_location=dist_util.dev() 137 | ) 138 | ema_params = self.mp_trainer.state_dict_to_master_params(state_dict) 139 | 140 | dist_util.sync_params(ema_params) 141 | return ema_params 142 | 143 | def _load_optimizer_state(self): 144 | main_checkpoint = find_resume_checkpoint() or self.resume_checkpoint 145 | opt_checkpoint = bf.join( 146 | bf.dirname(main_checkpoint), f"opt{self.resume_step:06}.pt" 147 | ) 148 | if bf.exists(opt_checkpoint): 149 | logger.log(f"loading optimizer state from checkpoint: {opt_checkpoint}") 150 | state_dict = dist_util.load_state_dict( 151 | opt_checkpoint, map_location=dist_util.dev() 152 | ) 153 | self.opt.load_state_dict(state_dict) 154 | 155 | def run_loop(self): 156 | logger.log("calculating training time...") 157 | start = time.time() 158 | while ( 159 | not self.lr_anneal_steps 160 | or self.step + self.resume_step < self.lr_anneal_steps 161 | ): 162 | batch, cond = next(self.data) 163 | self.run_step(batch, cond) 164 | if self.step % self.log_interval == 0: 165 | logger.dumpkvs() 166 | if self.step % self.save_interval == 0: 167 | end = time.time() 168 | logger.log(f"training time upto iter-{self.step}: {end - start:.5f} sec") 169 | self.save() 170 | # Run for a finite amount of time in integration tests. 171 | if os.environ.get("DIFFUSION_TRAINING_TEST", "") and self.step > 0: 172 | return 173 | self.step += 1 174 | # Save the last checkpoint if it wasn't already saved. 175 | if (self.step - 1) % self.save_interval != 0: 176 | self.save() 177 | 178 | def run_step(self, batch, cond): 179 | self.forward_backward(batch, cond) 180 | took_step = self.mp_trainer.optimize(self.opt) 181 | if took_step: 182 | self._update_ema() 183 | self._anneal_lr() 184 | self.log_step() 185 | 186 | def forward_backward(self, batch, cond): 187 | self.mp_trainer.zero_grad() 188 | for i in range(0, batch.shape[0], self.microbatch): 189 | micro = batch[i : i + self.microbatch].to(dist_util.dev()) 190 | micro_cond = { 191 | k: v[i : i + self.microbatch].to(dist_util.dev()) 192 | for k, v in cond.items() 193 | } 194 | last_batch = (i + self.microbatch) >= batch.shape[0] 195 | t, weights = self.schedule_sampler.sample(micro.shape[0], dist_util.dev()) 196 | 197 | compute_losses = functools.partial( 198 | self.diffusion.training_losses, 199 | self.ddp_model, 200 | micro, 201 | t, 202 | model_kwargs=micro_cond, 203 | ) 204 | 205 | if last_batch or not self.use_ddp: 206 | losses = compute_losses() 207 | else: 208 | with self.ddp_model.no_sync(): 209 | losses = compute_losses() 210 | 211 | if isinstance(self.schedule_sampler, LossAwareSampler): 212 | self.schedule_sampler.update_with_local_losses( 213 | t, losses["loss"].detach() 214 | ) 215 | 216 | loss = (losses["loss"] * weights).mean() 217 | log_loss_dict( 218 | self.diffusion, t, {k: v * weights for k, v in losses.items()} 219 | ) 220 | self.mp_trainer.backward(loss) 221 | 222 | def _update_ema(self): 223 | for rate, params in zip(self.ema_rate, self.ema_params): 224 | update_ema(params, self.mp_trainer.master_params, rate=rate) 225 | 226 | def _anneal_lr(self): 227 | if not self.lr_anneal_steps: 228 | return 229 | frac_done = (self.step + self.resume_step) / self.lr_anneal_steps 230 | lr = self.lr * (1 - frac_done) 231 | for param_group in self.opt.param_groups: 232 | param_group["lr"] = lr 233 | 234 | def log_step(self): 235 | logger.logkv("step", self.step + self.resume_step) 236 | logger.logkv("samples", (self.step + self.resume_step + 1) * self.global_batch) 237 | 238 | def save(self): 239 | def save_checkpoint(rate, params): 240 | state_dict = self.mp_trainer.master_params_to_state_dict(params) 241 | if dist.get_rank() == 0: 242 | logger.log(f"saving model {rate}...") 243 | if not rate: 244 | filename = f"model{(self.step+self.resume_step):06d}.pt" 245 | else: 246 | filename = f"ema_{rate}_{(self.step+self.resume_step):06d}.pt" 247 | with bf.BlobFile(bf.join(get_blob_logdir(), filename), "wb") as f: 248 | th.save(state_dict, f) 249 | 250 | save_checkpoint(0, self.mp_trainer.master_params) 251 | for rate, params in zip(self.ema_rate, self.ema_params): 252 | save_checkpoint(rate, params) 253 | 254 | if dist.get_rank() == 0: 255 | with bf.BlobFile( 256 | bf.join(get_blob_logdir(), f"opt{(self.step+self.resume_step):06d}.pt"), 257 | "wb", 258 | ) as f: 259 | th.save(self.opt.state_dict(), f) 260 | 261 | dist.barrier() 262 | 263 | 264 | def parse_resume_step_from_filename(filename): 265 | """ 266 | Parse filenames of the form path/to/modelNNNNNN.pt, where NNNNNN is the 267 | checkpoint's number of steps. 268 | """ 269 | split = filename.split("model") 270 | if len(split) < 2: 271 | return 0 272 | split1 = split[-1].split(".")[0] 273 | try: 274 | return int(split1) 275 | except ValueError: 276 | return 0 277 | 278 | 279 | def get_blob_logdir(): 280 | # You can change this to be a separate path to save checkpoints to 281 | # a blobstore or some external drive. 282 | return logger.get_dir() 283 | 284 | 285 | def find_resume_checkpoint(): 286 | # On your infrastructure, you may want to override this to automatically 287 | # discover the latest checkpoint on your blob storage, etc. 288 | return None 289 | 290 | 291 | def find_ema_checkpoint(main_checkpoint, step, rate): 292 | if main_checkpoint is None: 293 | return None 294 | filename = f"ema_{rate}_{(step):06d}.pt" 295 | path = bf.join(bf.dirname(main_checkpoint), filename) 296 | if bf.exists(path): 297 | return path 298 | return None 299 | 300 | 301 | def log_loss_dict(diffusion, ts, losses): 302 | for key, values in losses.items(): 303 | logger.logkv_mean(key, values.mean().item()) 304 | # Log the quantiles (four quartiles, in particular). 305 | for sub_t, sub_loss in zip(ts.cpu().numpy(), values.detach().cpu().numpy()): 306 | quartile = int(4 * sub_t / diffusion.num_timesteps) 307 | logger.logkv_mean(f"{key}_q{quartile}", sub_loss) 308 | -------------------------------------------------------------------------------- /guided_diffusion/script_util.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import inspect 3 | 4 | from . import gaussian_diffusion as gd 5 | from .respace import SpacedDiffusion, space_timesteps 6 | from .unet import SuperResModel, UNetModel, EncoderUNetModel 7 | 8 | NUM_CLASSES = 1000 9 | 10 | 11 | def diffusion_defaults(): 12 | """ 13 | Defaults for image and classifier training. 14 | """ 15 | return dict( 16 | learn_sigma=False, 17 | diffusion_steps=1000, 18 | noise_schedule="linear", 19 | timestep_respacing="", 20 | use_kl=False, 21 | predict_xstart=False, 22 | rescale_timesteps=False, 23 | rescale_learned_sigmas=False, 24 | ) 25 | 26 | 27 | def classifier_defaults(): 28 | """ 29 | Defaults for classifier models. 30 | """ 31 | return dict( 32 | image_size=64, 33 | classifier_use_fp16=False, 34 | classifier_width=128, 35 | classifier_depth=2, 36 | classifier_attention_resolutions="32,16,8", # 16 37 | classifier_use_scale_shift_norm=True, # False 38 | classifier_resblock_updown=True, # False 39 | classifier_pool="attention", 40 | ) 41 | 42 | 43 | def model_and_diffusion_defaults(): 44 | """ 45 | Defaults for image training. 46 | """ 47 | res = dict( 48 | image_size=64, 49 | num_channels=128, 50 | num_res_blocks=2, 51 | num_heads=4, 52 | num_heads_upsample=-1, 53 | num_head_channels=-1, 54 | attention_resolutions="16,8", 55 | channel_mult="", 56 | dropout=0.0, 57 | class_cond=False, 58 | use_checkpoint=False, 59 | use_scale_shift_norm=True, 60 | resblock_updown=False, 61 | use_fp16=False, 62 | use_new_attention_order=False, 63 | gray_imgs=False, ##### modified by sum: for admitting grayscale images ##### 64 | ) 65 | res.update(diffusion_defaults()) 66 | return res 67 | 68 | 69 | def classifier_and_diffusion_defaults(): 70 | res = classifier_defaults() 71 | res.update(diffusion_defaults()) 72 | return res 73 | 74 | 75 | def create_model_and_diffusion( 76 | image_size, 77 | class_cond, 78 | learn_sigma, 79 | num_channels, 80 | num_res_blocks, 81 | channel_mult, 82 | num_heads, 83 | num_head_channels, 84 | num_heads_upsample, 85 | attention_resolutions, 86 | dropout, 87 | diffusion_steps, 88 | noise_schedule, 89 | timestep_respacing, 90 | use_kl, 91 | predict_xstart, 92 | rescale_timesteps, 93 | rescale_learned_sigmas, 94 | use_checkpoint, 95 | use_scale_shift_norm, 96 | resblock_updown, 97 | use_fp16, 98 | use_new_attention_order, 99 | gray_imgs, ##### modified by sum: for admitting grayscale images ##### 100 | ): 101 | model = create_model( 102 | image_size, 103 | num_channels, 104 | num_res_blocks, 105 | channel_mult=channel_mult, 106 | learn_sigma=learn_sigma, 107 | class_cond=class_cond, 108 | use_checkpoint=use_checkpoint, 109 | attention_resolutions=attention_resolutions, 110 | num_heads=num_heads, 111 | num_head_channels=num_head_channels, 112 | num_heads_upsample=num_heads_upsample, 113 | use_scale_shift_norm=use_scale_shift_norm, 114 | dropout=dropout, 115 | resblock_updown=resblock_updown, 116 | use_fp16=use_fp16, 117 | use_new_attention_order=use_new_attention_order, 118 | gray_imgs=gray_imgs, ##### modified by sum: for admitting grayscale images ##### 119 | ) 120 | diffusion = create_gaussian_diffusion( 121 | steps=diffusion_steps, 122 | learn_sigma=learn_sigma, 123 | noise_schedule=noise_schedule, 124 | use_kl=use_kl, 125 | predict_xstart=predict_xstart, 126 | rescale_timesteps=rescale_timesteps, 127 | rescale_learned_sigmas=rescale_learned_sigmas, 128 | timestep_respacing=timestep_respacing, 129 | ) 130 | return model, diffusion 131 | 132 | 133 | def create_model( 134 | image_size, 135 | num_channels, 136 | num_res_blocks, 137 | channel_mult="", 138 | learn_sigma=False, 139 | class_cond=False, 140 | use_checkpoint=False, 141 | attention_resolutions="16", 142 | num_heads=1, 143 | num_head_channels=-1, 144 | num_heads_upsample=-1, 145 | use_scale_shift_norm=False, 146 | dropout=0, 147 | resblock_updown=False, 148 | use_fp16=False, 149 | use_new_attention_order=False, 150 | gray_imgs=False ##### modified by sum: for admitting grayscale images ##### 151 | ): 152 | if channel_mult == "": 153 | if image_size == 512: 154 | channel_mult = (0.5, 1, 1, 2, 2, 4, 4) 155 | elif image_size == 256: 156 | channel_mult = (1, 1, 2, 2, 4, 4) 157 | elif image_size == 128: 158 | channel_mult = (1, 1, 2, 3, 4) 159 | elif image_size == 64: 160 | channel_mult = (1, 2, 3, 4) 161 | else: 162 | raise ValueError(f"unsupported image size: {image_size}") 163 | else: 164 | channel_mult = tuple(int(ch_mult) for ch_mult in channel_mult.split(",")) 165 | 166 | attention_ds = [] 167 | for res in attention_resolutions.split(","): 168 | attention_ds.append(image_size // int(res)) 169 | 170 | return UNetModel( 171 | image_size=image_size, 172 | in_channels=(3 if not gray_imgs else 1), ##### modified by sum: for admitting grayscale images ##### 173 | model_channels=num_channels, 174 | out_channels=(3 if not learn_sigma else 6) if not gray_imgs else (1 if not learn_sigma else 2), ##### modified by sum: for admitting grayscale images ##### 175 | num_res_blocks=num_res_blocks, 176 | attention_resolutions=tuple(attention_ds), 177 | dropout=dropout, 178 | channel_mult=channel_mult, 179 | num_classes=(NUM_CLASSES if class_cond else None), 180 | use_checkpoint=use_checkpoint, 181 | use_fp16=use_fp16, 182 | num_heads=num_heads, 183 | num_head_channels=num_head_channels, 184 | num_heads_upsample=num_heads_upsample, 185 | use_scale_shift_norm=use_scale_shift_norm, 186 | resblock_updown=resblock_updown, 187 | use_new_attention_order=use_new_attention_order, 188 | ) 189 | 190 | 191 | def create_classifier_and_diffusion( 192 | image_size, 193 | classifier_use_fp16, 194 | classifier_width, 195 | classifier_depth, 196 | classifier_attention_resolutions, 197 | classifier_use_scale_shift_norm, 198 | classifier_resblock_updown, 199 | classifier_pool, 200 | learn_sigma, 201 | diffusion_steps, 202 | noise_schedule, 203 | timestep_respacing, 204 | use_kl, 205 | predict_xstart, 206 | rescale_timesteps, 207 | rescale_learned_sigmas, 208 | ): 209 | classifier = create_classifier( 210 | image_size, 211 | classifier_use_fp16, 212 | classifier_width, 213 | classifier_depth, 214 | classifier_attention_resolutions, 215 | classifier_use_scale_shift_norm, 216 | classifier_resblock_updown, 217 | classifier_pool, 218 | ) 219 | diffusion = create_gaussian_diffusion( 220 | steps=diffusion_steps, 221 | learn_sigma=learn_sigma, 222 | noise_schedule=noise_schedule, 223 | use_kl=use_kl, 224 | predict_xstart=predict_xstart, 225 | rescale_timesteps=rescale_timesteps, 226 | rescale_learned_sigmas=rescale_learned_sigmas, 227 | timestep_respacing=timestep_respacing, 228 | ) 229 | return classifier, diffusion 230 | 231 | 232 | def create_classifier( 233 | image_size, 234 | classifier_use_fp16, 235 | classifier_width, 236 | classifier_depth, 237 | classifier_attention_resolutions, 238 | classifier_use_scale_shift_norm, 239 | classifier_resblock_updown, 240 | classifier_pool, 241 | ): 242 | if image_size == 512: 243 | channel_mult = (0.5, 1, 1, 2, 2, 4, 4) 244 | elif image_size == 256: 245 | channel_mult = (1, 1, 2, 2, 4, 4) 246 | elif image_size == 128: 247 | channel_mult = (1, 1, 2, 3, 4) 248 | elif image_size == 64: 249 | channel_mult = (1, 2, 3, 4) 250 | else: 251 | raise ValueError(f"unsupported image size: {image_size}") 252 | 253 | attention_ds = [] 254 | for res in classifier_attention_resolutions.split(","): 255 | attention_ds.append(image_size // int(res)) 256 | 257 | return EncoderUNetModel( 258 | image_size=image_size, 259 | in_channels=3, 260 | model_channels=classifier_width, 261 | out_channels=1000, 262 | num_res_blocks=classifier_depth, 263 | attention_resolutions=tuple(attention_ds), 264 | channel_mult=channel_mult, 265 | use_fp16=classifier_use_fp16, 266 | num_head_channels=64, 267 | use_scale_shift_norm=classifier_use_scale_shift_norm, 268 | resblock_updown=classifier_resblock_updown, 269 | pool=classifier_pool, 270 | ) 271 | 272 | 273 | def sr_model_and_diffusion_defaults(): 274 | res = model_and_diffusion_defaults() 275 | res["large_size"] = 256 276 | res["small_size"] = 64 277 | arg_names = inspect.getfullargspec(sr_create_model_and_diffusion)[0] 278 | for k in res.copy().keys(): 279 | if k not in arg_names: 280 | del res[k] 281 | return res 282 | 283 | 284 | def sr_create_model_and_diffusion( 285 | large_size, 286 | small_size, 287 | class_cond, 288 | learn_sigma, 289 | num_channels, 290 | num_res_blocks, 291 | num_heads, 292 | num_head_channels, 293 | num_heads_upsample, 294 | attention_resolutions, 295 | dropout, 296 | diffusion_steps, 297 | noise_schedule, 298 | timestep_respacing, 299 | use_kl, 300 | predict_xstart, 301 | rescale_timesteps, 302 | rescale_learned_sigmas, 303 | use_checkpoint, 304 | use_scale_shift_norm, 305 | resblock_updown, 306 | use_fp16, 307 | ): 308 | model = sr_create_model( 309 | large_size, 310 | small_size, 311 | num_channels, 312 | num_res_blocks, 313 | learn_sigma=learn_sigma, 314 | class_cond=class_cond, 315 | use_checkpoint=use_checkpoint, 316 | attention_resolutions=attention_resolutions, 317 | num_heads=num_heads, 318 | num_head_channels=num_head_channels, 319 | num_heads_upsample=num_heads_upsample, 320 | use_scale_shift_norm=use_scale_shift_norm, 321 | dropout=dropout, 322 | resblock_updown=resblock_updown, 323 | use_fp16=use_fp16, 324 | ) 325 | diffusion = create_gaussian_diffusion( 326 | steps=diffusion_steps, 327 | learn_sigma=learn_sigma, 328 | noise_schedule=noise_schedule, 329 | use_kl=use_kl, 330 | predict_xstart=predict_xstart, 331 | rescale_timesteps=rescale_timesteps, 332 | rescale_learned_sigmas=rescale_learned_sigmas, 333 | timestep_respacing=timestep_respacing, 334 | ) 335 | return model, diffusion 336 | 337 | 338 | def sr_create_model( 339 | large_size, 340 | small_size, 341 | num_channels, 342 | num_res_blocks, 343 | learn_sigma, 344 | class_cond, 345 | use_checkpoint, 346 | attention_resolutions, 347 | num_heads, 348 | num_head_channels, 349 | num_heads_upsample, 350 | use_scale_shift_norm, 351 | dropout, 352 | resblock_updown, 353 | use_fp16, 354 | ): 355 | _ = small_size # hack to prevent unused variable 356 | 357 | if large_size == 512: 358 | channel_mult = (1, 1, 2, 2, 4, 4) 359 | elif large_size == 256: 360 | channel_mult = (1, 1, 2, 2, 4, 4) 361 | elif large_size == 64: 362 | channel_mult = (1, 2, 3, 4) 363 | else: 364 | raise ValueError(f"unsupported large size: {large_size}") 365 | 366 | attention_ds = [] 367 | for res in attention_resolutions.split(","): 368 | attention_ds.append(large_size // int(res)) 369 | 370 | return SuperResModel( 371 | image_size=large_size, 372 | in_channels=3, 373 | model_channels=num_channels, 374 | out_channels=(3 if not learn_sigma else 6), 375 | num_res_blocks=num_res_blocks, 376 | attention_resolutions=tuple(attention_ds), 377 | dropout=dropout, 378 | channel_mult=channel_mult, 379 | num_classes=(NUM_CLASSES if class_cond else None), 380 | use_checkpoint=use_checkpoint, 381 | num_heads=num_heads, 382 | num_head_channels=num_head_channels, 383 | num_heads_upsample=num_heads_upsample, 384 | use_scale_shift_norm=use_scale_shift_norm, 385 | resblock_updown=resblock_updown, 386 | use_fp16=use_fp16, 387 | ) 388 | 389 | 390 | def create_gaussian_diffusion( 391 | *, 392 | steps=1000, 393 | learn_sigma=False, 394 | sigma_small=False, 395 | noise_schedule="linear", 396 | use_kl=False, 397 | predict_xstart=False, 398 | rescale_timesteps=False, 399 | rescale_learned_sigmas=False, 400 | timestep_respacing="", 401 | ): 402 | betas = gd.get_named_beta_schedule(noise_schedule, steps) 403 | if use_kl: 404 | loss_type = gd.LossType.RESCALED_KL 405 | elif rescale_learned_sigmas: 406 | loss_type = gd.LossType.RESCALED_MSE 407 | else: 408 | loss_type = gd.LossType.MSE 409 | if not timestep_respacing: 410 | timestep_respacing = [steps] 411 | return SpacedDiffusion( 412 | use_timesteps=space_timesteps(steps, timestep_respacing), 413 | betas=betas, 414 | model_mean_type=( 415 | gd.ModelMeanType.EPSILON if not predict_xstart else gd.ModelMeanType.START_X 416 | ), 417 | model_var_type=( 418 | ( 419 | gd.ModelVarType.FIXED_LARGE 420 | if not sigma_small 421 | else gd.ModelVarType.FIXED_SMALL 422 | ) 423 | if not learn_sigma 424 | else gd.ModelVarType.LEARNED_RANGE 425 | ), 426 | loss_type=loss_type, 427 | rescale_timesteps=rescale_timesteps, 428 | ) 429 | 430 | 431 | def add_dict_to_argparser(parser, default_dict): 432 | for k, v in default_dict.items(): 433 | v_type = type(v) 434 | if v is None: 435 | v_type = str 436 | elif isinstance(v, bool): 437 | v_type = str2bool 438 | parser.add_argument(f"--{k}", default=v, type=v_type) 439 | 440 | 441 | def args_to_dict(args, keys): 442 | return {k: getattr(args, k) for k in keys} 443 | 444 | 445 | def str2bool(v): 446 | """ 447 | https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse 448 | """ 449 | if isinstance(v, bool): 450 | return v 451 | if v.lower() in ("yes", "true", "t", "y", "1"): 452 | return True 453 | elif v.lower() in ("no", "false", "f", "n", "0"): 454 | return False 455 | else: 456 | raise argparse.ArgumentTypeError("boolean value expected") 457 | -------------------------------------------------------------------------------- /baseline_algorithms/[adjoint_optimization] GA_waveguide.py: -------------------------------------------------------------------------------- 1 | import meep as mp 2 | import meep.adjoint as mpa 3 | 4 | from autograd import numpy as npa 5 | from autograd import tensor_jacobian_product, grad 6 | import matplotlib.pyplot as plt 7 | import nlopt 8 | import torch 9 | import numpy as np 10 | 11 | from scipy.ndimage import label, binary_dilation 12 | import os 13 | from skimage.measure import euler_number 14 | from skimage.measure import label as label_ 15 | 16 | import wandb 17 | wandb.init(project = "adjoint_waveguide_GD", settings=wandb.Settings(silent=True)) 18 | 19 | directory = 'adjoint/new' 20 | 21 | import os 22 | if not os.path.exists(directory): 23 | os.makedirs(directory) 24 | 25 | plt.rcParams["font.family"] = "sans-serif" 26 | plt.rcParams["font.sans-serif"] = ["DejaVu Sans"] 27 | 28 | plt.rcParams["figure.figsize"] = (3.5,3.5) 29 | 30 | 31 | 32 | params = { 33 | 'axes.labelsize':9, # label 폰트 크기 34 | 'axes.titlesize':9, # 타이틀 폰트 크기 35 | 'xtick.labelsize':9, # x 축 tick label 폰트 크기 36 | 'ytick.labelsize':9, # y 축 tick label 폰트 크기 37 | 'xtick.direction': 'in', # 눈금 표시 방향 (in, out, inout) 38 | 'ytick.direction': 'in', # 눈금 표시 방향 (in, out, inout) 39 | 'lines.markersize': 3, # 마커 사이즈 40 | 'axes.titlepad': 6, # 타이틀과 그래프 사이의 간격 41 | 'axes.labelpad': 4, # 축 label과 그래프 사이의 간격 42 | 'font.size': 9, # font 크기 43 | #'font.sans-serif': 'Arial', # font 설정 44 | 'figure.dpi': 300, # 해상도, vector그래픽의 경우 dpi에 상관없이 깔끔하게 출력됨 45 | 'figure.autolayout': True, # 레이아웃 자동 설정 (그래프의 모든 요소가 figure 내부에 들어가도록 설정) 46 | 'xtick.top': True, # 그래프 위쪽 x축 눈금 표시 47 | 'ytick.right': True, # 그래프 오른쪽 y축 눈금 표시 48 | 'xtick.major.size': 2, # x축 눈금의 길이 49 | 'ytick.major.size': 2, # y축 눈금의 길이 50 | } 51 | 52 | 53 | 54 | # Function to delete every island with size 1 55 | def delete_islands_with_size_1(array): 56 | # Label the connected components 57 | labeled_array, num_features = label(array) 58 | 59 | # Identify the features (islands) with size 1 60 | islands_to_delete = [i for i in range(1, num_features + 1) if np.sum(labeled_array == i) == 1] 61 | 62 | # Delete these islands 63 | for island in islands_to_delete: 64 | array[labeled_array == island] = 0 65 | 66 | if array.ndim > 2: 67 | array = np.squeeze(array, axis=2) 68 | else: 69 | pass 70 | return array 71 | 72 | 73 | # Function to find the minimum feature size and its location 74 | def find_minimum_feature_size(array): 75 | # Label the connected components 76 | labeled_array, num_features = label(array) 77 | 78 | # Get the sizes of all features 79 | feature_sizes = [(i, np.sum(labeled_array == i)) for i in range(1, num_features + 1)] 80 | 81 | # Find the minimum feature size and its label 82 | min_feature = min(feature_sizes, key=lambda x: x[1]) if feature_sizes else (0, 0) 83 | 84 | return min_feature[1], min_feature[0], labeled_array 85 | 86 | 87 | 88 | def highlight_minimum_island(array, min_label, labeled_array, title): 89 | # Find the coordinates of the minimum island 90 | coords = np.argwhere(labeled_array == min_label) 91 | if coords.size == 0: 92 | return array 93 | 94 | # Create an empty mask for the boundary 95 | boundary_mask = np.zeros_like(array) 96 | 97 | # Fill the mask with the minimum island 98 | boundary_mask[labeled_array == min_label] = 1 99 | 100 | # Dilate the mask and subtract the original to get the boundary 101 | dilated_mask = binary_dilation(boundary_mask) 102 | boundary = dilated_mask - boundary_mask 103 | 104 | # Create a plot 105 | fig, ax = plt.subplots() 106 | ax.imshow(1-array, cmap='gray') # black: 1, while: 0 107 | ax.imshow(boundary, cmap='Reds', alpha=0.5) # Overlay the red boundary with higher alpha 108 | ax.set_title(title) 109 | 110 | # Adjust boundary line thickness and make it more vivid 111 | boundary_indices = np.argwhere(boundary) 112 | for y, x in boundary_indices: 113 | rect = plt.Rectangle((x - 0.5, y - 0.5), 1, 1, edgecolor='red', facecolor='red', linewidth=1) 114 | ax.add_patch(rect) 115 | 116 | return plt 117 | 118 | mp.verbosity(0) 119 | Si = mp.Medium(index=3.4) 120 | SiO2 = mp.Medium(index=1.44) 121 | 122 | 123 | resolution = 21 124 | 125 | Sx = 10 126 | Sy = 10 127 | cell_size = mp.Vector3(Sx, Sy) 128 | 129 | pml_layers = [mp.PML(2.0)] 130 | 131 | fcen = 1 / 1.55 132 | width = 0.2 133 | fwidth = width * fcen 134 | source_center = [-2.7, 0, 0] 135 | source_size = mp.Vector3(0, 2, 0) 136 | kpoint = mp.Vector3(1, 0, 0) 137 | src = mp.GaussianSource(frequency=fcen, fwidth=fwidth) 138 | source = [ 139 | mp.EigenModeSource( 140 | src, 141 | eig_band=1, 142 | direction=mp.NO_DIRECTION, 143 | eig_kpoint=kpoint, 144 | size=source_size, 145 | center=source_center, 146 | ) 147 | ] 148 | 149 | design_region_resolution = 21 150 | Nx = 64#design_region_resolution +1 151 | Ny = 64#design_region_resolution +1 152 | 153 | design_variables = mp.MaterialGrid(mp.Vector3(Nx, Ny), SiO2, Si, grid_type="U_MEAN") 154 | design_region = mpa.DesignRegion( 155 | design_variables, volume=mp.Volume(center=mp.Vector3(), size=mp.Vector3(3, 3, 0)) 156 | ) 157 | 158 | """ 159 | mp.Block( 160 | center=mp.Vector3(y=Sy / 4), material=Si, size=mp.Vector3(1, Sy / 2, 0) 161 | ), # vertical waveguide 162 | """ 163 | geometry = [ 164 | mp.Block( 165 | center=mp.Vector3(x=-Sx / 4), material=Si, size=mp.Vector3(Sx / 2, 1, 0) 166 | ), # horizontal waveguide 167 | #mp.Block( 168 | # center=mp.Vector3(x=Sx / 4), material=Si, size=mp.Vector3(Sx / 2, 1, 0) 169 | #), # vertical waveguide 170 | mp.Block( 171 | center=mp.Vector3(y=Sy / 4), material=Si, size=mp.Vector3(1, Sy / 2, 0) 172 | ), # vertical waveguide 173 | mp.Block( 174 | center=design_region.center, size=design_region.size, material=design_variables 175 | ), # design region 176 | # mp.Block(center=design_region.center, size=design_region.size, material=design_variables, 177 | # e1=mp.Vector3(x=-1).rotate(mp.Vector3(z=1), np.pi/2), e2=mp.Vector3(y=1).rotate(mp.Vector3(z=1), np.pi/2)) 178 | # 179 | # The commented lines above impose symmetry by overlapping design region with the same design variable. However, 180 | # currently there is an issue of doing that; We give an alternative approach to impose symmetry in later tutorials. 181 | # See https://github.com/NanoComp/meep/issues/1984 and https://github.com/NanoComp/meep/issues/2093 182 | ] 183 | 184 | sim = mp.Simulation( 185 | cell_size=cell_size, 186 | boundary_layers=pml_layers, 187 | geometry=geometry, 188 | sources=source, 189 | eps_averaging=True, 190 | subpixel_tol=1e-4, 191 | resolution=resolution, 192 | ) 193 | 194 | TE_front = mpa.EigenmodeCoefficient( 195 | sim, mp.Volume(center=mp.Vector3(2.5, 0, 0), size=mp.Vector3(y=2)), mode=1 196 | ) 197 | TE_top = mpa.EigenmodeCoefficient( 198 | sim, mp.Volume(center=mp.Vector3(0, 2.5, 0), size=mp.Vector3(x=2)), mode=1 199 | ) 200 | TE_bot = mpa.EigenmodeCoefficient( 201 | sim, mp.Volume(center=mp.Vector3(0, -2.5, 0), size=mp.Vector3(x=2)), mode=1, forward = False 202 | ) 203 | TE_0 = mpa.EigenmodeCoefficient( 204 | sim, mp.Volume(center=mp.Vector3(-2.5, 0, 0), size=mp.Vector3(y=2)), mode=1 205 | ) 206 | #ob_list = [TE0, TE_top] 207 | ob_list = [TE_0, TE_top] 208 | 209 | def J (source, top): 210 | return npa.abs(top / source) ** 2 211 | 212 | ''' 213 | def conic_filter(x, R, Nx, Ny): 214 | x_new = np.zeros_like(x) 215 | 216 | for i in range(Nx): 217 | for j in range(Ny): 218 | weights_sum = 0 219 | values_sum = 0 220 | for m in range(Nx): 221 | for n in range(Ny): 222 | d = np.sqrt((i - m)**2 + (j - n)**2) 223 | if d <= R: 224 | weight = max(0, 1 - d/R) 225 | weights_sum += weight 226 | values_sum += weight * x[m * Ny + n] 227 | 228 | x_new[i * Ny + j] = values_sum / weights_sum 229 | 230 | return x_new 231 | ''' 232 | 233 | 234 | minimum_length = 0.895 # minimum length scale (microns) 235 | eta_i = 0.5 # blueprint (or intermediate) design field thresholding point (between 0 and 1) 236 | eta_e = 0.65 # erosion design field thresholding point (between 0 and 1) 237 | eta_d = 1 - eta_e # dilation design field thresholding point (between 0 and 1) 238 | filter_radius = minimum_length#mpa.get_conic_radius_from_eta_e(minimum_length, eta_e) 239 | design_region_width = 3 240 | design_region_height = 3 241 | 242 | wandb.config.update({"minimum length": minimum_length, "resolution": resolution, 243 | "Sx": Sx, "Sy": Sy, "Nx": Nx, "Ny": Ny, 244 | "eta_i": eta_i, "eta_e": eta_e, "eta_d": eta_d 245 | }) 246 | 247 | def mapping (x, eta, beta): 248 | # filter 씌우기 (블러처리느낌) 249 | 250 | x = mpa.conic_filter( 251 | x, 252 | filter_radius, 253 | design_region_width, 254 | design_region_height, 255 | design_region_resolution 256 | ) 257 | x = x.flatten () 258 | # 출력값 .1~ 1으로 제한 259 | 260 | print('filtered_field type: ', type(x)) 261 | print("filtered_field shape:", getattr(x, 'shape', 'No shape attribute')) 262 | 263 | if not isinstance(x, np.ndarray): 264 | raise TypeError("Expected filtered_field to be an ndarray, got {}".format(type(filtered_field))) 265 | 266 | beta = float(beta) 267 | eta = float(eta) 268 | projected_field = mpa.tanh_projection(x, beta, eta) 269 | # interpolate to actual materials 270 | return projected_field.flatten () 271 | 272 | opt = mpa.OptimizationProblem( 273 | simulation=sim, 274 | objective_functions=J, 275 | objective_arguments=ob_list, 276 | design_regions=[design_region], 277 | fcen=fcen, 278 | df=0, 279 | nf=1, 280 | ) 281 | 282 | #struct_random = np.random.choice([0, 1], size=(Nx*Ny)) 283 | 284 | import random 285 | x = 0.5 * np.ones((Nx * Ny,)) 286 | #x0 = np.random.randint(2, size= Nx*Ny) 287 | #x0 = np.random.rand((Nx*Ny)) 288 | opt.update_design([x]) 289 | 290 | evaluation_history = [] 291 | sensitivity = [0] 292 | 293 | 294 | 295 | 296 | n = Nx*Ny 297 | LR = 1 298 | cur_beta = 2 299 | beta_scale = 2 300 | num_betas = 140 301 | 302 | 303 | wandb.run.name = "min"+str(minimum_length)+"_"+"eta"+str(eta_i)+"_"+"num_betas"+str(num_betas)+"_beta_scale"+str(beta_scale)+"_LR"+str(LR) 304 | 305 | 306 | wandb.config.update({ 307 | 'beta_scale': beta_scale, 308 | 'num_betas': num_betas}) 309 | 310 | x = np.ones(n)*0.5 311 | for iters in range(num_betas): 312 | x[x>1] = 1 313 | x[x<0] = 0 314 | #print('cur beta: ', cur_beta) 315 | print('x: ', x) 316 | f0, dJ_du = opt([mapping(x, eta_i, cur_beta)]) 317 | adjoint_gradient = dJ_du 318 | reshaped_gradients = adjoint_gradient.reshape(adjoint_gradient.shape[0], -1) 319 | 320 | # Computing the norm of each reshaped row 321 | norms = np.linalg.norm(reshaped_gradients, axis=1) 322 | 323 | # Calculating the mean of the norms 324 | adjgrad_norm = norms.mean() 325 | print(adjoint_gradient) 326 | print(type(adjoint_gradient)) 327 | 328 | 329 | x_image = x.reshape(Nx, Ny,1) 330 | print(x_image.shape) 331 | 332 | learning_rate = LR/ adjgrad_norm 333 | #learning_rate = LR 334 | wandb.log({ 335 | 'fom ': f0, 336 | 'adjgrad norm': adjgrad_norm, 337 | 'learning rate': learning_rate, 338 | "generated": [wandb.Image(x_image, caption=f"Iteration {iters}")] 339 | 340 | }), 341 | x = x + learning_rate * adjoint_gradient 342 | # what if learning_rate = learning_rate / 343 | cur_beta = cur_beta * beta_scale 344 | 345 | 346 | 347 | ## final 348 | # Final thresholding of the design : must be deleted!! 349 | #x [x>0.5] = 1 350 | #x [x==0.5] = 1 351 | #x [x<0.5] = 0 352 | 353 | x_final = mapping(x, eta_i, mp.inf) 354 | 355 | f0, dJ_du = opt([x_final]) 356 | sensitivity[0] = dJ_du 357 | x_image = x.reshape(Nx, Ny,1) 358 | 359 | wandb.log({ 360 | 'final_fom ': np.real(f0), 361 | "generated_final": [wandb.Image(x_image, caption=f"generated_final")] 362 | 363 | }), 364 | 365 | fig = plt.figure(figsize=(20,20)) 366 | plt.imshow((np.squeeze(np.abs(sensitivity[0].reshape(Nx, Ny)))))#(np.rot90(np.squeeze(np.abs(sensitivity[0].reshape(Nx, Ny))))); 367 | plt.xlabel("x") 368 | plt.ylabel("y") 369 | plt.savefig("adjoint/new/sensitivity.png") 370 | cmap = plt.cm.get_cmap() 371 | colormapping = plt.cm.ScalarMappable(cmap=cmap) 372 | cbar = fig.colorbar(colormapping, ax=plt.gca()) 373 | wandb.log({"sensitivity": plt}) 374 | 375 | 376 | 377 | x_final = x 378 | 379 | x_final = x_final.reshape(Nx, Ny) 380 | 381 | island_deleted = delete_islands_with_size_1(x_final.copy()) 382 | min_feature_size, min_feature_label, labeled_array = find_minimum_feature_size(island_deleted) 383 | 384 | fom_island, _ = opt([ 385 | island_deleted.flatten()]) 386 | print('min feature size: ', min_feature_size) 387 | print('min feature label: ', min_feature_label) 388 | print('labeled array: ', labeled_array) 389 | print('island_deleted: ', island_deleted.shape) 390 | 391 | #plt1 = highlight_minimum_island(island_deleted, min_feature_label, labeled_array, title='Minimum Feature 1s') 392 | 393 | wandb.log({ 394 | 'fom (final, island deleted1)': fom_island, 395 | "generated_final, island deleted1": [wandb.Image(1-island_deleted, caption=f"generated_final_island_deleted")], 396 | # "highlighted, island deleted1": [wandb.Image(plt1, caption=f"generated_final_island_deleted_highlghted")] 397 | }) 398 | array_converted_processed = delete_islands_with_size_1(1-island_deleted.copy()) 399 | array_original = 1- array_converted_processed 400 | 401 | min_feature_size2, min_feature_label2, labeled_array2 = find_minimum_feature_size(array_converted_processed) 402 | plt2 = highlight_minimum_island(array_original, min_feature_label2, labeled_array2, title='Minimum Feature 0s') 403 | euler_number_original = euler_number(array_original, connectivity=1) 404 | _, num_islands1 = label_(island_deleted, connectivity=1, return_num=True) 405 | num_holes1 = num_islands1 - euler_number_original 406 | 407 | euler_number_converted = euler_number(array_converted_processed, connectivity=1) 408 | _, num_islands2 = label_(array_converted_processed, connectivity=1, return_num=True) 409 | num_holes2 = num_islands2 - euler_number_converted 410 | 411 | 412 | fom_island2, _ = opt([ 413 | array_original.flatten()] 414 | ) 415 | wandb.log({ 416 | "highlighted, island deleted2": [wandb.Image(plt2, caption=f"generated_final_island_deleted_highlghted")], 417 | 418 | }) 419 | plt3 = highlight_minimum_island(array_original, min_feature_label, labeled_array, title='Minimum Feature 1s') 420 | 421 | 422 | wandb.log({ 423 | 'fom (final, island deleted2)': fom_island2, 424 | "generated_final, island deleted2": [wandb.Image(1-array_original, caption=f"generated_final_island_deleted")], 425 | "highlighted, final": [wandb.Image(plt3, caption=f"generated_final_island_deleted_highlghted")] 426 | }) 427 | 428 | print("array 0,0: ", array_original[0,0]) 429 | 430 | 431 | wandb.log({ 432 | 'mfs (original)': min_feature_size, 433 | 'mfs (converted)': min_feature_size2, 434 | 'number of islands (original)': num_islands1, 435 | 'number of islands (converted)':num_islands2, 436 | 'number of holes (original)' : num_holes1, 437 | 'number of holes (converted)': num_holes2, 438 | 'euler_number with 1 connectivity (original)': euler_number_original, 439 | 'euler_number with 1 connectivity (converted)': euler_number_converted, 440 | }) 441 | 442 | 443 | print("DONE") 444 | 445 | -------------------------------------------------------------------------------- /guided_diffusion/logger.py: -------------------------------------------------------------------------------- 1 | """ 2 | Logger copied from OpenAI baselines to avoid extra RL-based dependencies: 3 | https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/logger.py 4 | """ 5 | 6 | import os 7 | import sys 8 | import shutil 9 | import os.path as osp 10 | import json 11 | import time 12 | import datetime 13 | import tempfile 14 | import warnings 15 | from collections import defaultdict 16 | from contextlib import contextmanager 17 | 18 | DEBUG = 10 19 | INFO = 20 20 | WARN = 30 21 | ERROR = 40 22 | 23 | DISABLED = 50 24 | 25 | 26 | class KVWriter(object): 27 | def writekvs(self, kvs): 28 | raise NotImplementedError 29 | 30 | 31 | class SeqWriter(object): 32 | def writeseq(self, seq): 33 | raise NotImplementedError 34 | 35 | 36 | class HumanOutputFormat(KVWriter, SeqWriter): 37 | def __init__(self, filename_or_file): 38 | if isinstance(filename_or_file, str): 39 | self.file = open(filename_or_file, "wt") 40 | self.own_file = True 41 | else: 42 | assert hasattr(filename_or_file, "read"), ( 43 | "expected file or str, got %s" % filename_or_file 44 | ) 45 | self.file = filename_or_file 46 | self.own_file = False 47 | 48 | def writekvs(self, kvs): 49 | # Create strings for printing 50 | key2str = {} 51 | for (key, val) in sorted(kvs.items()): 52 | if hasattr(val, "__float__"): 53 | valstr = "%-8.3g" % val 54 | else: 55 | valstr = str(val) 56 | key2str[self._truncate(key)] = self._truncate(valstr) 57 | 58 | # Find max widths 59 | if len(key2str) == 0: 60 | print("WARNING: tried to write empty key-value dict") 61 | return 62 | else: 63 | keywidth = max(map(len, key2str.keys())) 64 | valwidth = max(map(len, key2str.values())) 65 | 66 | # Write out the data 67 | dashes = "-" * (keywidth + valwidth + 7) 68 | lines = [dashes] 69 | for (key, val) in sorted(key2str.items(), key=lambda kv: kv[0].lower()): 70 | lines.append( 71 | "| %s%s | %s%s |" 72 | % (key, " " * (keywidth - len(key)), val, " " * (valwidth - len(val))) 73 | ) 74 | lines.append(dashes) 75 | self.file.write("\n".join(lines) + "\n") 76 | 77 | # Flush the output to the file 78 | self.file.flush() 79 | 80 | def _truncate(self, s): 81 | maxlen = 30 82 | return s[: maxlen - 3] + "..." if len(s) > maxlen else s 83 | 84 | def writeseq(self, seq): 85 | seq = list(seq) 86 | for (i, elem) in enumerate(seq): 87 | self.file.write(elem) 88 | if i < len(seq) - 1: # add space unless this is the last one 89 | self.file.write(" ") 90 | self.file.write("\n") 91 | self.file.flush() 92 | 93 | def close(self): 94 | if self.own_file: 95 | self.file.close() 96 | 97 | 98 | class JSONOutputFormat(KVWriter): 99 | def __init__(self, filename): 100 | self.file = open(filename, "wt") 101 | 102 | def writekvs(self, kvs): 103 | for k, v in sorted(kvs.items()): 104 | if hasattr(v, "dtype"): 105 | kvs[k] = float(v) 106 | self.file.write(json.dumps(kvs) + "\n") 107 | self.file.flush() 108 | 109 | def close(self): 110 | self.file.close() 111 | 112 | 113 | class CSVOutputFormat(KVWriter): 114 | def __init__(self, filename): 115 | self.file = open(filename, "w+t") 116 | self.keys = [] 117 | self.sep = "," 118 | 119 | def writekvs(self, kvs): 120 | # Add our current row to the history 121 | extra_keys = list(kvs.keys() - self.keys) 122 | extra_keys.sort() 123 | if extra_keys: 124 | self.keys.extend(extra_keys) 125 | self.file.seek(0) 126 | lines = self.file.readlines() 127 | self.file.seek(0) 128 | for (i, k) in enumerate(self.keys): 129 | if i > 0: 130 | self.file.write(",") 131 | self.file.write(k) 132 | self.file.write("\n") 133 | for line in lines[1:]: 134 | self.file.write(line[:-1]) 135 | self.file.write(self.sep * len(extra_keys)) 136 | self.file.write("\n") 137 | for (i, k) in enumerate(self.keys): 138 | if i > 0: 139 | self.file.write(",") 140 | v = kvs.get(k) 141 | if v is not None: 142 | self.file.write(str(v)) 143 | self.file.write("\n") 144 | self.file.flush() 145 | 146 | def close(self): 147 | self.file.close() 148 | 149 | 150 | class TensorBoardOutputFormat(KVWriter): 151 | """ 152 | Dumps key/value pairs into TensorBoard's numeric format. 153 | """ 154 | 155 | def __init__(self, dir): 156 | os.makedirs(dir, exist_ok=True) 157 | self.dir = dir 158 | self.step = 1 159 | prefix = "events" 160 | path = osp.join(osp.abspath(dir), prefix) 161 | import tensorflow as tf 162 | from tensorflow.python import pywrap_tensorflow 163 | from tensorflow.core.util import event_pb2 164 | from tensorflow.python.util import compat 165 | 166 | self.tf = tf 167 | self.event_pb2 = event_pb2 168 | self.pywrap_tensorflow = pywrap_tensorflow 169 | self.writer = pywrap_tensorflow.EventsWriter(compat.as_bytes(path)) 170 | 171 | def writekvs(self, kvs): 172 | def summary_val(k, v): 173 | kwargs = {"tag": k, "simple_value": float(v)} 174 | return self.tf.Summary.Value(**kwargs) 175 | 176 | summary = self.tf.Summary(value=[summary_val(k, v) for k, v in kvs.items()]) 177 | event = self.event_pb2.Event(wall_time=time.time(), summary=summary) 178 | event.step = ( 179 | self.step 180 | ) # is there any reason why you'd want to specify the step? 181 | self.writer.WriteEvent(event) 182 | self.writer.Flush() 183 | self.step += 1 184 | 185 | def close(self): 186 | if self.writer: 187 | self.writer.Close() 188 | self.writer = None 189 | 190 | 191 | def make_output_format(format, ev_dir, log_suffix=""): 192 | os.makedirs(ev_dir, exist_ok=True) 193 | if format == "stdout": 194 | return HumanOutputFormat(sys.stdout) 195 | elif format == "log": 196 | return HumanOutputFormat(osp.join(ev_dir, "log%s.txt" % log_suffix)) 197 | elif format == "json": 198 | return JSONOutputFormat(osp.join(ev_dir, "progress%s.json" % log_suffix)) 199 | elif format == "csv": 200 | return CSVOutputFormat(osp.join(ev_dir, "progress%s.csv" % log_suffix)) 201 | elif format == "tensorboard": 202 | return TensorBoardOutputFormat(osp.join(ev_dir, "tb%s" % log_suffix)) 203 | else: 204 | raise ValueError("Unknown format specified: %s" % (format,)) 205 | 206 | 207 | # ================================================================ 208 | # API 209 | # ================================================================ 210 | 211 | 212 | def logkv(key, val): 213 | """ 214 | Log a value of some diagnostic 215 | Call this once for each diagnostic quantity, each iteration 216 | If called many times, last value will be used. 217 | """ 218 | get_current().logkv(key, val) 219 | 220 | 221 | def logkv_mean(key, val): 222 | """ 223 | The same as logkv(), but if called many times, values averaged. 224 | """ 225 | get_current().logkv_mean(key, val) 226 | 227 | 228 | def logkvs(d): 229 | """ 230 | Log a dictionary of key-value pairs 231 | """ 232 | for (k, v) in d.items(): 233 | logkv(k, v) 234 | 235 | 236 | def dumpkvs(): 237 | """ 238 | Write all of the diagnostics from the current iteration 239 | """ 240 | return get_current().dumpkvs() 241 | 242 | 243 | def getkvs(): 244 | return get_current().name2val 245 | 246 | 247 | def log(*args, level=INFO): 248 | """ 249 | Write the sequence of args, with no separators, to the console and output files (if you've configured an output file). 250 | """ 251 | get_current().log(*args, level=level) 252 | 253 | 254 | def debug(*args): 255 | log(*args, level=DEBUG) 256 | 257 | 258 | def info(*args): 259 | log(*args, level=INFO) 260 | 261 | 262 | def warn(*args): 263 | log(*args, level=WARN) 264 | 265 | 266 | def error(*args): 267 | log(*args, level=ERROR) 268 | 269 | 270 | def set_level(level): 271 | """ 272 | Set logging threshold on current logger. 273 | """ 274 | get_current().set_level(level) 275 | 276 | 277 | def set_comm(comm): 278 | get_current().set_comm(comm) 279 | 280 | 281 | def get_dir(): 282 | """ 283 | Get directory that log files are being written to. 284 | will be None if there is no output directory (i.e., if you didn't call start) 285 | """ 286 | return get_current().get_dir() 287 | 288 | 289 | record_tabular = logkv 290 | dump_tabular = dumpkvs 291 | 292 | 293 | @contextmanager 294 | def profile_kv(scopename): 295 | logkey = "wait_" + scopename 296 | tstart = time.time() 297 | try: 298 | yield 299 | finally: 300 | get_current().name2val[logkey] += time.time() - tstart 301 | 302 | 303 | def profile(n): 304 | """ 305 | Usage: 306 | @profile("my_func") 307 | def my_func(): code 308 | """ 309 | 310 | def decorator_with_name(func): 311 | def func_wrapper(*args, **kwargs): 312 | with profile_kv(n): 313 | return func(*args, **kwargs) 314 | 315 | return func_wrapper 316 | 317 | return decorator_with_name 318 | 319 | 320 | # ================================================================ 321 | # Backend 322 | # ================================================================ 323 | 324 | 325 | def get_current(): 326 | if Logger.CURRENT is None: 327 | _configure_default_logger() 328 | 329 | return Logger.CURRENT 330 | 331 | 332 | class Logger(object): 333 | DEFAULT = None # A logger with no output files. (See right below class definition) 334 | # So that you can still log to the terminal without setting up any output files 335 | CURRENT = None # Current logger being used by the free functions above 336 | 337 | def __init__(self, dir, output_formats, comm=None): 338 | self.name2val = defaultdict(float) # values this iteration 339 | self.name2cnt = defaultdict(int) 340 | self.level = INFO 341 | self.dir = dir 342 | self.output_formats = output_formats 343 | self.comm = comm 344 | 345 | # Logging API, forwarded 346 | # ---------------------------------------- 347 | def logkv(self, key, val): 348 | self.name2val[key] = val 349 | 350 | def logkv_mean(self, key, val): 351 | oldval, cnt = self.name2val[key], self.name2cnt[key] 352 | self.name2val[key] = oldval * cnt / (cnt + 1) + val / (cnt + 1) 353 | self.name2cnt[key] = cnt + 1 354 | 355 | def dumpkvs(self): 356 | if self.comm is None: 357 | d = self.name2val 358 | else: 359 | d = mpi_weighted_mean( 360 | self.comm, 361 | { 362 | name: (val, self.name2cnt.get(name, 1)) 363 | for (name, val) in self.name2val.items() 364 | }, 365 | ) 366 | if self.comm.rank != 0: 367 | d["dummy"] = 1 # so we don't get a warning about empty dict 368 | out = d.copy() # Return the dict for unit testing purposes 369 | for fmt in self.output_formats: 370 | if isinstance(fmt, KVWriter): 371 | fmt.writekvs(d) 372 | self.name2val.clear() 373 | self.name2cnt.clear() 374 | return out 375 | 376 | def log(self, *args, level=INFO): 377 | if self.level <= level: 378 | self._do_log(args) 379 | 380 | # Configuration 381 | # ---------------------------------------- 382 | def set_level(self, level): 383 | self.level = level 384 | 385 | def set_comm(self, comm): 386 | self.comm = comm 387 | 388 | def get_dir(self): 389 | return self.dir 390 | 391 | def close(self): 392 | for fmt in self.output_formats: 393 | fmt.close() 394 | 395 | # Misc 396 | # ---------------------------------------- 397 | def _do_log(self, args): 398 | for fmt in self.output_formats: 399 | if isinstance(fmt, SeqWriter): 400 | fmt.writeseq(map(str, args)) 401 | 402 | 403 | def get_rank_without_mpi_import(): 404 | # check environment variables here instead of importing mpi4py 405 | # to avoid calling MPI_Init() when this module is imported 406 | for varname in ["PMI_RANK", "OMPI_COMM_WORLD_RANK"]: 407 | if varname in os.environ: 408 | return int(os.environ[varname]) 409 | return 0 410 | 411 | 412 | def mpi_weighted_mean(comm, local_name2valcount): 413 | """ 414 | Copied from: https://github.com/openai/baselines/blob/ea25b9e8b234e6ee1bca43083f8f3cf974143998/baselines/common/mpi_util.py#L110 415 | Perform a weighted average over dicts that are each on a different node 416 | Input: local_name2valcount: dict mapping key -> (value, count) 417 | Returns: key -> mean 418 | """ 419 | all_name2valcount = comm.gather(local_name2valcount) 420 | if comm.rank == 0: 421 | name2sum = defaultdict(float) 422 | name2count = defaultdict(float) 423 | for n2vc in all_name2valcount: 424 | for (name, (val, count)) in n2vc.items(): 425 | try: 426 | val = float(val) 427 | except ValueError: 428 | if comm.rank == 0: 429 | warnings.warn( 430 | "WARNING: tried to compute mean on non-float {}={}".format( 431 | name, val 432 | ) 433 | ) 434 | else: 435 | name2sum[name] += val * count 436 | name2count[name] += count 437 | return {name: name2sum[name] / name2count[name] for name in name2sum} 438 | else: 439 | return {} 440 | 441 | 442 | def configure(dir=None, format_strs=None, comm=None, log_suffix=""): 443 | """ 444 | If comm is provided, average all numerical stats across that comm 445 | """ 446 | if dir is None: 447 | dir = os.getenv("OPENAI_LOGDIR") 448 | if dir is None: 449 | dir = osp.join( 450 | tempfile.gettempdir(), 451 | datetime.datetime.now().strftime("openai-%Y-%m-%d-%H-%M-%S-%f"), 452 | ) 453 | assert isinstance(dir, str) 454 | dir = os.path.expanduser(dir) 455 | os.makedirs(os.path.expanduser(dir), exist_ok=True) 456 | 457 | rank = get_rank_without_mpi_import() 458 | if rank > 0: 459 | log_suffix = log_suffix + "-rank%03i" % rank 460 | 461 | if format_strs is None: 462 | if rank == 0: 463 | format_strs = os.getenv("OPENAI_LOG_FORMAT", "stdout,log,csv").split(",") 464 | else: 465 | format_strs = os.getenv("OPENAI_LOG_FORMAT_MPI", "log").split(",") 466 | format_strs = filter(None, format_strs) 467 | output_formats = [make_output_format(f, dir, log_suffix) for f in format_strs] 468 | 469 | Logger.CURRENT = Logger(dir=dir, output_formats=output_formats, comm=comm) 470 | if output_formats: 471 | log("Logging to %s" % dir) 472 | 473 | 474 | def _configure_default_logger(): 475 | configure() 476 | Logger.DEFAULT = Logger.CURRENT 477 | 478 | 479 | def reset(): 480 | if Logger.CURRENT is not Logger.DEFAULT: 481 | Logger.CURRENT.close() 482 | Logger.CURRENT = Logger.DEFAULT 483 | log("Reset logger") 484 | 485 | 486 | @contextmanager 487 | def scoped_configure(dir=None, format_strs=None, comm=None): 488 | prevlogger = Logger.CURRENT 489 | configure(dir=dir, format_strs=format_strs, comm=comm) 490 | try: 491 | yield 492 | finally: 493 | Logger.CURRENT.close() 494 | Logger.CURRENT = prevlogger 495 | 496 | -------------------------------------------------------------------------------- /simulation.py: -------------------------------------------------------------------------------- 1 | import matplotlib.pyplot as plt 2 | import pickle 3 | import os 4 | 5 | # Configure global matplotlib settings 6 | plt.rcParams["font.family"] = "sans-serif" 7 | plt.rcParams["font.sans-serif"] = ["DejaVu Sans"] 8 | plt.rcParams["figure.figsize"] = (3.5, 3.5) 9 | 10 | params = { 11 | 'axes.labelsize': 12, # font size for axis labels 12 | 'axes.titlesize': 12, # font size for plot titles 13 | 'xtick.labelsize': 10, # font size for x-axis tick labels 14 | 'ytick.labelsize': 10, # font size for y-axis tick labels 15 | 'xtick.direction': 'in', # tick direction (in, out, inout) 16 | 'ytick.direction': 'in', # tick direction (in, out, inout) 17 | 'lines.markersize': 3, # marker size 18 | 'axes.titlepad': 6, # padding between title and plot 19 | 'axes.labelpad': 4, # padding between axis label and plot 20 | 'font.size': 12, # global font size 21 | # 'font.sans-serif': 'Arial', # Uncomment to specify a particular sans-serif font 22 | 'figure.dpi': 300, # DPI setting (for raster graphics) 23 | 'figure.autolayout': True, # automatically adjust layout to fit all elements 24 | 'xtick.top': True, # show top x-axis ticks 25 | 'ytick.right': True, # show right y-axis ticks 26 | 'xtick.major.size': 2, # x-axis major tick length 27 | 'ytick.major.size': 2, # y-axis major tick length 28 | } 29 | 30 | # Function to save lists to a pickle file 31 | def save_lists(new_red_list, new_blue_list, new_green_list): 32 | """ 33 | Appends new red, blue, and green data to existing lists 34 | and saves them in 'lists.pkl' as a pickle file. 35 | """ 36 | if os.path.exists('lists.pkl'): 37 | with open('lists.pkl', 'rb') as f: 38 | red_list, blue_list, green_list = pickle.load(f) 39 | else: 40 | red_list, blue_list, green_list = [], [], [] 41 | 42 | red_list += new_red_list 43 | blue_list += new_blue_list 44 | green_list += new_green_list 45 | 46 | with open('lists.pkl', 'wb') as f: 47 | pickle.dump((red_list, blue_list, green_list), f) 48 | 49 | def CIS_sim(struct_np, t, exp_name, prop_dir='', 50 | save_inter=False, interval=10, flag_last=False): 51 | """ 52 | CIS Simulation function: 53 | 54 | - Takes an array 'struct_np' which is used as the design region's binary (0/1) pattern. 55 | - Sets up a MEEP simulation for a 3x3 design region on top of an SiO2 substrate block 56 | with side blocks of SiPD, aiming to capture transmissions of red, green, and blue 57 | frequencies. 58 | - Saves the summed intensities of the fields in separate lists (red, green, blue), 59 | which are stored using 'save_lists()'. 60 | - Returns the figure of merit (sum of the intensities) and its gradient w.r.t. 61 | the design variables. 62 | """ 63 | 64 | import numpy as np 65 | import meep as mp 66 | import meep.adjoint as mpa 67 | import autograd.numpy as npa 68 | import matplotlib.pyplot as plt 69 | import wandb # If not using wandb, remove or comment out. 70 | 71 | # Ensure structure is clipped between 0 and 1 72 | struct_np[struct_np > 1] = 1 73 | struct_np[struct_np < 0] = 0 74 | 75 | red_list = [] 76 | blue_list = [] 77 | green_list = [] 78 | 79 | mp.verbosity(0) 80 | 81 | # Define materials 82 | Air = mp.Medium(index=1.0) 83 | SiN = mp.Medium(epsilon=4) 84 | SiO2 = mp.Medium(epsilon=2.1) 85 | SiPD = mp.Medium(epsilon=5) 86 | 87 | # Basic simulation parameters 88 | um_scale = 1 89 | resolution = 21 90 | 91 | # Design region dimensions 92 | design_region_width = 3 93 | design_region_height = 3 94 | 95 | # Additional geometry 96 | gapop = 0 97 | air_gap = 0 98 | dti = 0.4 99 | subpixelsize = design_region_width / 3 - dti 100 | if gapop == 1: 101 | air_gap = dti / 2 102 | PDsize = 2 103 | Lpml = 0.5 104 | Sourcespace = 2 105 | 106 | # Cell size 107 | Sx = design_region_width 108 | Sy = PDsize + design_region_height + Sourcespace + Lpml 109 | cell_size = mp.Vector3(Sx, Sy) 110 | 111 | # PML settings 112 | pml_layers = [mp.PML(thickness=Lpml, direction=mp.Y)] 113 | 114 | # Frequency ranges 115 | wavelengths = np.linspace(0.40 * um_scale, 0.70 * um_scale, 31) 116 | frequencies = 1 / wavelengths 117 | nf = len(frequencies) 118 | 119 | # Source frequency setup (red, green, blue) 120 | width = 0.4 121 | fcen_red = 1 / (0.65 * um_scale) 122 | fwidth_red = fcen_red * width 123 | 124 | fcen_green = 1 / (0.55 * um_scale) 125 | fwidth_green = fcen_green * width 126 | 127 | fcen_blue = 1 / (0.45 * um_scale) 128 | fwidth_blue = fcen_blue * width 129 | 130 | src_0 = mp.GaussianSource(frequency=fcen_red, fwidth=fwidth_red, is_integrated=True) 131 | src_1 = mp.GaussianSource(frequency=fcen_green, fwidth=fwidth_green, is_integrated=True) 132 | src_2 = mp.GaussianSource(frequency=fcen_blue, fwidth=fwidth_blue, is_integrated=True) 133 | 134 | source_center = mp.Vector3(0, Sy / 2 - Lpml - Sourcespace / 2, 0) 135 | source_size = mp.Vector3(Sx, 0, 0) 136 | source = [ 137 | mp.Source(src_0, component=mp.Ez, size=source_size, center=source_center), 138 | mp.Source(src_1, component=mp.Ez, size=source_size, center=source_center), 139 | mp.Source(src_2, component=mp.Ez, size=source_size, center=source_center), 140 | ] 141 | 142 | # Grid dimensions 143 | Nx = 64 144 | Ny = 64 145 | 146 | # Create a MaterialGrid for the design region 147 | design_variables = mp.MaterialGrid( 148 | mp.Vector3(Nx, Ny), 149 | SiO2, 150 | SiN, 151 | grid_type="U_MEAN" 152 | ) 153 | 154 | design_region = mpa.DesignRegion( 155 | design_variables, 156 | volume=mp.Volume( 157 | center=mp.Vector3(0, -Sy/2 + PDsize + design_region_height / 2, 0), 158 | size=mp.Vector3(design_region_width, design_region_height, 0), 159 | ), 160 | ) 161 | 162 | # Geometry 163 | geometry = [ 164 | mp.Block( 165 | center=design_region.center, 166 | size=design_region.size, 167 | material=design_variables 168 | ), 169 | mp.Block( 170 | center=mp.Vector3(0, -Sy/2 + PDsize/2, 0), 171 | size=mp.Vector3(Sx, PDsize, 0), 172 | material=SiO2 173 | ), 174 | # Three blocks of SiPD 175 | mp.Block( 176 | center=mp.Vector3(-design_region_width/3, -Sy/2 + PDsize/2, 0), 177 | size=mp.Vector3(subpixelsize, PDsize, 0), 178 | material=SiPD 179 | ), 180 | mp.Block( 181 | center=mp.Vector3(0, -Sy/2 + PDsize/2, 0), 182 | size=mp.Vector3(subpixelsize, PDsize, 0), 183 | material=SiPD 184 | ), 185 | mp.Block( 186 | center=mp.Vector3(design_region_width/3, -Sy/2 + PDsize/2, 0), 187 | size=mp.Vector3(subpixelsize, PDsize, 0), 188 | material=SiPD 189 | ), 190 | ] 191 | 192 | # Set up simulation 193 | sim = mp.Simulation( 194 | cell_size=cell_size, 195 | boundary_layers=pml_layers, 196 | geometry=geometry, 197 | sources=source, 198 | default_material=Air, 199 | resolution=resolution, 200 | k_point=mp.Vector3(0, 0, 0), 201 | ) 202 | 203 | # Define monitor positions and sizes 204 | monitor_position_0 = mp.Vector3(-design_region_width/3, -Sy/2 + PDsize - 0.5/resolution) 205 | monitor_size_0 = mp.Vector3(subpixelsize, 0) 206 | 207 | monitor_position_1 = mp.Vector3(0, -Sy/2 + PDsize - 0.5/resolution) 208 | monitor_size_1 = mp.Vector3(subpixelsize, 0) 209 | 210 | monitor_position_2 = mp.Vector3(design_region_width/3, -Sy/2 + PDsize - 0.5/resolution) 211 | monitor_size_2 = mp.Vector3(subpixelsize, 0) 212 | 213 | # Create FourierFields for each position 214 | FourierFields_0 = mpa.FourierFields(sim, mp.Volume(center=monitor_position_0, size=monitor_size_0), mp.Ez, yee_grid=True) 215 | FourierFields_1 = mpa.FourierFields(sim, mp.Volume(center=monitor_position_1, size=monitor_size_1), mp.Ez, yee_grid=True) 216 | FourierFields_2 = mpa.FourierFields(sim, mp.Volume(center=monitor_position_2, size=monitor_size_2), mp.Ez, yee_grid=True) 217 | 218 | ob_list = [FourierFields_0, FourierFields_1, FourierFields_2] 219 | 220 | def J_0(fields_0, fields_1, fields_2): 221 | """ 222 | Objective function that sums intensities in three frequency bands 223 | corresponding to red, green, and blue. 224 | """ 225 | red = npa.sum(npa.abs(fields_0[21:30, :]) ** 2) 226 | green = npa.sum(npa.abs(fields_1[11:20, :]) ** 2) 227 | blue = npa.sum(npa.abs(fields_2[1:10, :]) ** 2) 228 | 229 | # Ensure we handle AutoDiff types vs. raw floats 230 | red_ = red._value if hasattr(red, '_value') else red 231 | green_ = green._value if hasattr(green, '_value') else green 232 | blue_ = blue._value if hasattr(blue, '_value') else blue 233 | 234 | # Append and save 235 | red_list.append(red_) 236 | blue_list.append(blue_) 237 | green_list.append(green_) 238 | save_lists(red_list, blue_list, green_list) 239 | 240 | return blue + green + red 241 | 242 | # Define optimization problem 243 | opt = mpa.OptimizationProblem( 244 | simulation=sim, 245 | objective_functions=[J_0], 246 | objective_arguments=ob_list, 247 | design_regions=[design_region], 248 | frequencies=frequencies, 249 | decay_by=1e-3, # field decay threshold 250 | ) 251 | 252 | # Flatten and update design 253 | flattened_array = struct_np.flatten() 254 | opt.update_design([flattened_array]) 255 | 256 | # Compute figure of merit and gradient 257 | fom, dJ_du = opt([flattened_array]) 258 | g = npa.sum(dJ_du, axis=1) 259 | 260 | return fom, g 261 | 262 | def waveguide_sim(struct_np, t, exp_name, prop_dir='top', 263 | save_inter=False, interval=1, flag_last=False): 264 | """ 265 | Waveguide simulation function: 266 | 267 | - Takes a 2D 'struct_np' array (clipped to 0/1). 268 | - Places it as a design region within a simple waveguide geometry in MEEP. 269 | - 'prop_dir' specifies which waveguide we are measuring transmission into 270 | ('top', 'bottom', or 'front'). 271 | - Returns the figure of merit and its gradient. 272 | """ 273 | 274 | import numpy as np 275 | import meep as mp 276 | import meep.adjoint as mpa 277 | import autograd.numpy as npa 278 | import matplotlib.pyplot as plt 279 | import wandb # If not using wandb, remove or comment out. 280 | 281 | mp.verbosity(0) 282 | 283 | # Clip design values 284 | struct_np[struct_np > 1] = 1 285 | struct_np[struct_np < 0] = 0 286 | struct_np = np.squeeze(struct_np) 287 | 288 | # Define materials 289 | Si = mp.Medium(index=3.4) 290 | SiO2 = mp.Medium(index=1.44) 291 | 292 | resolution = 21 293 | 294 | # Cell size 295 | Sx = 10 296 | Sy = 10 297 | cell_size = mp.Vector3(Sx, Sy) 298 | pml_layers = [mp.PML(2.0)] 299 | 300 | # Frequency setup 301 | fcen = 1 / 1.55 302 | width = 0.2 303 | fwidth = fcen * width 304 | 305 | # Source setup 306 | source_center = mp.Vector3(-2.7, 0, 0) 307 | source_size = mp.Vector3(0, 2, 0) 308 | kpoint = mp.Vector3(1, 0, 0) 309 | src = mp.GaussianSource(frequency=fcen, fwidth=fwidth) 310 | source = [ 311 | mp.EigenModeSource( 312 | src, 313 | eig_band=1, 314 | direction=mp.NO_DIRECTION, 315 | eig_kpoint=kpoint, 316 | size=source_size, 317 | center=source_center, 318 | ) 319 | ] 320 | 321 | design_region_resolution = 21 322 | Nx = 64 323 | Ny = 64 324 | 325 | # Define the design region 326 | design_variables = mp.MaterialGrid( 327 | mp.Vector3(Nx, Ny), 328 | SiO2, 329 | Si, 330 | grid_type="U_MEAN" 331 | ) 332 | design_region = mpa.DesignRegion( 333 | design_variables, 334 | volume=mp.Volume( 335 | center=mp.Vector3(), 336 | size=mp.Vector3(3, 3, 0) 337 | ) 338 | ) 339 | 340 | # Geometry depending on 'prop_dir' 341 | if prop_dir == 'top': 342 | geometry = [ 343 | # Horizontal waveguide on the left (origin) 344 | mp.Block( 345 | center=mp.Vector3(x=-Sx / 4), 346 | material=Si, 347 | size=mp.Vector3(Sx / 2, 1, 0) 348 | ), 349 | # Vertical waveguide on the top 350 | mp.Block( 351 | center=mp.Vector3(y=Sy / 4), 352 | material=Si, 353 | size=mp.Vector3(1, Sy / 2, 0) 354 | ), 355 | # Design region 356 | mp.Block( 357 | center=design_region.center, 358 | size=design_region.size, 359 | material=design_variables 360 | ), 361 | ] 362 | elif prop_dir == 'bottom': 363 | geometry = [ 364 | # Horizontal waveguide on the left (origin) 365 | mp.Block( 366 | center=mp.Vector3(x=-Sx / 4), 367 | material=Si, 368 | size=mp.Vector3(Sx / 2, 1, 0) 369 | ), 370 | # Vertical waveguide on the bottom 371 | mp.Block( 372 | center=mp.Vector3(y=-Sy / 4), 373 | material=Si, 374 | size=mp.Vector3(1, Sy / 2, 0) 375 | ), 376 | # Design region 377 | mp.Block( 378 | center=design_region.center, 379 | size=design_region.size, 380 | material=design_variables 381 | ), 382 | ] 383 | elif prop_dir == 'front': 384 | geometry = [ 385 | # Horizontal waveguide on the left 386 | mp.Block( 387 | center=mp.Vector3(x=-Sx / 4), 388 | material=Si, 389 | size=mp.Vector3(Sx / 2, 1, 0) 390 | ), 391 | # Horizontal waveguide on the right ("front") 392 | mp.Block( 393 | center=mp.Vector3(x=Sx / 4), 394 | material=Si, 395 | size=mp.Vector3(Sx / 2, 1, 0) 396 | ), 397 | # Design region 398 | mp.Block( 399 | center=design_region.center, 400 | size=design_region.size, 401 | material=design_variables 402 | ), 403 | ] 404 | 405 | # Create the simulation 406 | sim = mp.Simulation( 407 | cell_size=cell_size, 408 | boundary_layers=pml_layers, 409 | geometry=geometry, 410 | sources=source, 411 | eps_averaging=True, 412 | subpixel_tol=1e-4, 413 | resolution=resolution, 414 | ) 415 | 416 | # Define mode monitors 417 | TE_front = mpa.EigenmodeCoefficient( 418 | sim, mp.Volume(center=mp.Vector3(2.5, 0, 0), size=mp.Vector3(y=2)), mode=1 419 | ) 420 | TE_top = mpa.EigenmodeCoefficient( 421 | sim, mp.Volume(center=mp.Vector3(0, 2.5, 0), size=mp.Vector3(x=2)), mode=1 422 | ) 423 | TE_bottom = mpa.EigenmodeCoefficient( 424 | sim, mp.Volume(center=mp.Vector3(0, -2.5, 0), size=mp.Vector3(x=2)), mode=1, forward=False 425 | ) 426 | TE_o = mpa.EigenmodeCoefficient( 427 | sim, mp.Volume(center=mp.Vector3(-2.5, 0, 0), size=mp.Vector3(y=2)), mode=1 428 | ) 429 | 430 | ob_list = [TE_o] # Always include the source monitor 431 | if prop_dir == 'top': 432 | ob_list.append(TE_top) 433 | elif prop_dir == 'bottom': 434 | ob_list.append(TE_bottom) 435 | elif prop_dir == 'front': 436 | ob_list.append(TE_front) 437 | 438 | def J(source_coeff, target_coeff): 439 | """ 440 | Objective function: maximize transmission 441 | = |target_coeff / source_coeff|^2 442 | """ 443 | return npa.abs(target_coeff / source_coeff) ** 2 444 | 445 | opt = mpa.OptimizationProblem( 446 | simulation=sim, 447 | objective_functions=J, 448 | objective_arguments=ob_list, 449 | design_regions=[design_region], 450 | fcen=fcen, 451 | df=0, 452 | nf=1, 453 | ) 454 | 455 | # Flatten and update design 456 | flattened_array = struct_np.flatten() 457 | opt.update_design([flattened_array]) 458 | 459 | # Compute figure of merit (FOM) and gradient 460 | fom, g = opt([flattened_array]) 461 | fom = fom[0] # Single frequency -> single objective 462 | 463 | return fom, g 464 | -------------------------------------------------------------------------------- /baseline_algorithms/[adjoint_optimization] nlopt_waveguide.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import nlopt 4 | import wandb 5 | import torch 6 | import numpy as np 7 | import matplotlib.pyplot as plt 8 | 9 | import meep as mp 10 | import meep.adjoint as mpa 11 | from scipy.ndimage import label, binary_dilation 12 | from skimage.measure import euler_number 13 | from skimage.measure import label as label_ 14 | from autograd import numpy as npa 15 | from autograd import tensor_jacobian_product 16 | 17 | 18 | # Initialize Weights & Biases (WandB) 19 | wandb.init(project="adjoint_waveguide", settings=wandb.Settings(silent=True)) 20 | 21 | # Create a directory for saving results 22 | directory = 'adjoint/new' 23 | if not os.path.exists(directory): 24 | os.makedirs(directory) 25 | 26 | # Matplotlib configurations 27 | plt.rcParams["font.family"] = "sans-serif" 28 | plt.rcParams["font.sans-serif"] = ["DejaVu Sans"] 29 | plt.rcParams["figure.figsize"] = (3.5, 3.5) 30 | params = { 31 | 'axes.labelsize': 12, # label font size 32 | 'axes.titlesize': 12, # title font size 33 | 'xtick.labelsize': 10, # x-axis tick label font size 34 | 'ytick.labelsize': 10, # y-axis tick label font size 35 | 'xtick.direction': 'in', # tick marking direction 36 | 'ytick.direction': 'in', 37 | 'lines.markersize': 3, 38 | 'axes.titlepad': 6, 39 | 'axes.labelpad': 4, 40 | 'font.size': 12, 41 | 'figure.dpi': 300, 42 | 'figure.autolayout': True, 43 | 'xtick.top': True, 44 | 'ytick.right': True, 45 | 'xtick.major.size': 2, 46 | 'ytick.major.size': 2 47 | } 48 | plt.rcParams.update(params) 49 | 50 | 51 | ############################################################################### 52 | # Utility / Helper Functions # 53 | ############################################################################### 54 | 55 | def delete_islands_with_size_1(array): 56 | """ 57 | Function to delete every island (connected component) with size 1 in the array. 58 | Uses scipy.ndimage.label. 59 | """ 60 | labeled_array, num_features = label(array) 61 | islands_to_delete = [i for i in range(1, num_features + 1) 62 | if np.sum(labeled_array == i) == 1] 63 | for island in islands_to_delete: 64 | array[labeled_array == island] = 0 65 | 66 | # If 3D array, squeeze the last axis if it's size=1 67 | if array.ndim > 2: 68 | array = np.squeeze(array, axis=2) 69 | return array 70 | 71 | 72 | def find_minimum_feature_size(array): 73 | """ 74 | Find the minimum feature size and its label in a binary array. 75 | Returns (minimum_feature_size, label_of_that_feature, labeled_array). 76 | """ 77 | labeled_array, num_features = label(array) 78 | feature_sizes = [(i, np.sum(labeled_array == i)) 79 | for i in range(1, num_features + 1)] 80 | min_feature = min(feature_sizes, key=lambda x: x[1]) if feature_sizes else (0, 0) 81 | return min_feature[1], min_feature[0], labeled_array 82 | 83 | 84 | def highlight_minimum_island(array, min_label, labeled_array): 85 | """ 86 | Creates a matplotlib plot highlighting the boundary of the island 87 | with 'min_label' in the given 'array'. 88 | """ 89 | coords = np.argwhere(labeled_array == min_label) 90 | if coords.size == 0: 91 | return array 92 | 93 | boundary_mask = np.zeros_like(array) 94 | boundary_mask[labeled_array == min_label] = 1 95 | 96 | dilated_mask = binary_dilation(boundary_mask) 97 | boundary = dilated_mask - boundary_mask 98 | 99 | fig, ax = plt.subplots() 100 | ax.imshow(1 - array, cmap='gray') # black: 1, white: 0 101 | ax.imshow(boundary, cmap='Reds', alpha=0.5) 102 | 103 | # Enhance boundary appearance 104 | boundary_indices = np.argwhere(boundary) 105 | for y, x in boundary_indices: 106 | rect = plt.Rectangle((x - 0.5, y - 0.5), 1, 1, 107 | edgecolor='red', facecolor='red', linewidth=1) 108 | ax.add_patch(rect) 109 | return plt 110 | 111 | 112 | ############################################################################### 113 | # Meep and Adjoint Setup # 114 | ############################################################################### 115 | 116 | # Meep setup 117 | mp.verbosity(0) 118 | Si = mp.Medium(index=3.4) 119 | SiO2 = mp.Medium(index=1.44) 120 | 121 | resolution = 21 122 | Sx = 10 123 | Sy = 10 124 | cell_size = mp.Vector3(Sx, Sy) 125 | pml_layers = [mp.PML(2.0)] 126 | 127 | fcen = 1 / 1.55 128 | width = 0.2 129 | fwidth = width * fcen 130 | source_center = [-2.7, 0, 0] 131 | source_size = mp.Vector3(0, 2, 0) 132 | kpoint = mp.Vector3(1, 0, 0) 133 | 134 | src = mp.GaussianSource(frequency=fcen, fwidth=fwidth) 135 | source = [ 136 | mp.EigenModeSource( 137 | src, 138 | eig_band=1, 139 | direction=mp.NO_DIRECTION, 140 | eig_kpoint=kpoint, 141 | size=source_size, 142 | center=source_center, 143 | ) 144 | ] 145 | 146 | # Design region resolution, grid size, and MaterialGrid 147 | design_region_resolution = 21 148 | Nx = 64 149 | Ny = 64 150 | 151 | design_variables = mp.MaterialGrid( 152 | mp.Vector3(Nx, Ny), 153 | SiO2, 154 | Si, 155 | grid_type="U_MEAN" 156 | ) 157 | 158 | design_region = mpa.DesignRegion( 159 | design_variables, 160 | volume=mp.Volume(center=mp.Vector3(), size=mp.Vector3(3, 3, 0)) 161 | ) 162 | 163 | ############################################################################### 164 | # Geometry Definition # 165 | ############################################################################### 166 | geometry = [ 167 | mp.Block( 168 | center=mp.Vector3(x=-Sx / 4), 169 | material=Si, 170 | size=mp.Vector3(Sx / 2, 1, 0) 171 | ), # horizontal waveguide 172 | 173 | mp.Block( 174 | center=mp.Vector3(y=Sy / 4), 175 | material=Si, 176 | size=mp.Vector3(1, Sy / 2, 0) 177 | ), # vertical waveguide 178 | 179 | mp.Block( 180 | center=design_region.center, 181 | size=design_region.size, 182 | material=design_variables 183 | ), 184 | ] 185 | 186 | sim = mp.Simulation( 187 | cell_size=cell_size, 188 | boundary_layers=pml_layers, 189 | geometry=geometry, 190 | sources=source, 191 | eps_averaging=True, 192 | subpixel_tol=1e-4, 193 | resolution=resolution, 194 | ) 195 | 196 | ############################################################################### 197 | # Define Observables # 198 | ############################################################################### 199 | TE_front = mpa.EigenmodeCoefficient( 200 | sim, mp.Volume(center=mp.Vector3(2.5, 0, 0), size=mp.Vector3(y=2)), mode=1 201 | ) 202 | TE_top = mpa.EigenmodeCoefficient( 203 | sim, mp.Volume(center=mp.Vector3(0, 2.5, 0), size=mp.Vector3(x=2)), mode=1 204 | ) 205 | TE_bot = mpa.EigenmodeCoefficient( 206 | sim, mp.Volume(center=mp.Vector3(0, -2.5, 0), size=mp.Vector3(x=2)), mode=1, forward=False 207 | ) 208 | TE_0 = mpa.EigenmodeCoefficient( 209 | sim, mp.Volume(center=mp.Vector3(-2.5, 0, 0), size=mp.Vector3(y=2)), mode=1 210 | ) 211 | 212 | # We consider TE_0 (source side) and TE_top as the metrics for objective 213 | ob_list = [TE_0, TE_top] 214 | 215 | def J(source, top): 216 | """Objective function: maximize |top/source|^2.""" 217 | return npa.abs(top / source) ** 2 218 | 219 | ############################################################################### 220 | # Optimization Setup # 221 | ############################################################################### 222 | minimum_length = 0.895 223 | eta_i = 0.5 224 | eta_e = 0.75 225 | eta_d = 1 - eta_e 226 | filter_radius = minimum_length # You can also use mpa.get_conic_radius_from_eta_e(...) 227 | design_region_width = 3 228 | design_region_height = 3 229 | 230 | wandb.config.update({ 231 | "minimum length": minimum_length, 232 | "resolution": resolution, 233 | "Sx": Sx, 234 | "Sy": Sy, 235 | "Nx": Nx, 236 | "Ny": Ny, 237 | "eta_i": eta_i, 238 | "eta_e": eta_e, 239 | "eta_d": eta_d 240 | }) 241 | 242 | def mapping(x, eta, beta): 243 | """ 244 | Mapping function that: 245 | 1) Applies a conic filter (like a blur) to implement a minimum length scale. 246 | 2) Projects the values using tanh_projection. 247 | 3) Returns an array in [0,1]. 248 | """ 249 | x_filtered = mpa.conic_filter( 250 | x, 251 | filter_radius, 252 | design_region_width, 253 | design_region_height, 254 | design_region_resolution 255 | ) 256 | projected_field = mpa.tanh_projection(x_filtered, beta, eta) 257 | return projected_field.flatten() 258 | 259 | # Create the optimization problem 260 | opt = mpa.OptimizationProblem( 261 | simulation=sim, 262 | objective_functions=J, 263 | objective_arguments=ob_list, 264 | design_regions=[design_region], 265 | fcen=fcen, 266 | df=0, 267 | nf=1, 268 | ) 269 | 270 | # Initialize design variables 271 | x = 0.5 * np.ones((Nx * Ny,)) 272 | opt.update_design([x]) 273 | 274 | evaluation_history = [] 275 | sensitivity = [0] 276 | 277 | # Set up optimization algorithms 278 | algorithm1 = nlopt.LD_MMA 279 | algorithm2 = nlopt.LD_SLSQP 280 | algorithm_dict = { 281 | algorithm1: 'MMA', 282 | algorithm2: 'SLSQP', 283 | } 284 | 285 | def algorithm_name(alg): 286 | return algorithm_dict.get(alg, "Unknown Algorithm") 287 | 288 | algorithm = algorithm2 289 | 290 | n = Nx * Ny 291 | MAXEVAL = 250 292 | eta = 0.5 293 | k = 0 294 | R = 8 295 | 296 | def f(x, grad, beta): 297 | """ 298 | The callback function for the optimizer. 299 | Evaluates the objective, computes gradients, logs to WandB, 300 | and saves snapshots. 301 | """ 302 | global k 303 | 304 | # Map the design variable with current threshold parameters 305 | x_mapped = mapping(x, eta, beta) 306 | 307 | # Evaluate objective & gradient 308 | f0, dJ_du = opt([x_mapped]) 309 | adjoint_gradient = dJ_du 310 | reshaped_gradients = adjoint_gradient.reshape(adjoint_gradient.shape[0], -1) 311 | 312 | # Compute norm for logging 313 | norms = np.linalg.norm(reshaped_gradients, axis=1) 314 | adjgrad_norm = norms.mean() 315 | 316 | print('fom: ', f0) 317 | print('adjgrad_norm: ', adjgrad_norm) 318 | 319 | # Log images to WandB 320 | wandb.log({ 321 | "fom": f0, 322 | "adjgrad_norm": adjgrad_norm, 323 | "generated": [wandb.Image(1 - x_mapped.reshape(Nx, Ny), 324 | caption='step' + str(k) + '_fom' + str(f0)[1:5])] 325 | }, step=k) 326 | 327 | # f0 is an array of length 1, so take its first (only) element 328 | f0_scalar = f0[0] 329 | 330 | # If gradient is requested by the optimizer 331 | if grad.size > 0: 332 | grad[:] = tensor_jacobian_product(mapping, 0)(x, eta_i, beta, dJ_du) 333 | 334 | evaluation_history.append(np.real(f0_scalar)) 335 | print(evaluation_history) 336 | sensitivity[0] = dJ_du 337 | 338 | # Update design in meep-adjoint 339 | opt.update_design([x_mapped]) 340 | opt.plot2D() 341 | 342 | # Save figure and data 343 | filename_prefix = f'adjoint/new/{str(k).zfill(2)}_{str(f0_scalar)[:5]}' 344 | plt.savefig(filename_prefix + '.png') 345 | np.save(filename_prefix + '.npy', x) 346 | 347 | k += 1 348 | return np.real(f0_scalar) 349 | 350 | cur_beta = 2 351 | beta_scale = 2 352 | num_betas = 7 353 | update_factor = 20 354 | 355 | wandb.run.name = (f"min{minimum_length}_eta_i{eta_i}_" 356 | f"{algorithm_name(algorithm)}_num_betas{num_betas}") 357 | 358 | for iters in range(num_betas): 359 | print("current beta: ", cur_beta) 360 | 361 | solver = nlopt.opt(algorithm, n) 362 | solver.set_lower_bounds(0) 363 | solver.set_upper_bounds(1) 364 | print('cur_beta: ', cur_beta) 365 | 366 | # If beta is large enough, effectively no more projection smoothing 367 | if cur_beta >= 2 ** (num_betas + 1): 368 | solver.set_max_objective(lambda a, g: f(a, g, mp.inf)) 369 | solver.set_maxeval(1) 370 | else: 371 | solver.set_max_objective(lambda a, g: f(a, g, cur_beta)) 372 | solver.set_maxeval(update_factor) 373 | 374 | np.savetxt('adjoint/new/history.txt', evaluation_history) 375 | x[x > 1] = 1 376 | x[x < 0] = 0 377 | x = solver.optimize(x) 378 | 379 | cur_beta *= beta_scale 380 | wandb.log({'cur_beta': cur_beta}, step=k) 381 | wandb.config.update({ 382 | 'update_factor': update_factor, 383 | 'beta_scale': beta_scale, 384 | 'num_betas': num_betas 385 | }) 386 | wandb.config.update({'algorithm': algorithm_name(algorithm)}) 387 | 388 | # Final thresholding of the design : must be deleted!! 389 | #x[x >= 0.5] = 1 390 | #x[x < 0.5] = 0 391 | 392 | x_final = mapping(x, eta_i, mp.inf) 393 | f0, dJ_du = opt([x_final]) 394 | sensitivity[0] = dJ_du 395 | 396 | wandb.log({ 397 | 'final_fom': np.real(f0), 398 | "generated_final": [wandb.Image(1 - design_variables.weights.reshape(Nx, Ny))], 399 | }) 400 | 401 | ############################################################################### 402 | # Plot Optimization Progress # 403 | ############################################################################### 404 | plt.figure() 405 | plt.plot(np.array(evaluation_history), "o-") 406 | plt.grid(True) 407 | plt.xlabel("Iteration") 408 | plt.ylabel("efficiency") 409 | plt.savefig('adjoint/new/log.png') 410 | 411 | fig = plt.figure(figsize=(5, 5)) 412 | plt.imshow(np.squeeze(np.abs(sensitivity[0].reshape(Nx, Ny)))) 413 | plt.xlabel("x") 414 | plt.ylabel("y") 415 | plt.savefig("adjoint/new/sensitivity.png") 416 | 417 | cmap = plt.cm.get_cmap() 418 | colormapping = plt.cm.ScalarMappable(cmap=cmap) 419 | cbar = fig.colorbar(colormapping, ax=plt.gca()) 420 | wandb.log({"sensitivity": plt}) 421 | 422 | ############################################################################### 423 | # Post-Processing of the Design # 424 | ############################################################################### 425 | x_final = x_final.reshape(Nx, Ny) 426 | 427 | # 1) Delete small islands (size=1) 428 | island_deleted = delete_islands_with_size_1(x_final.copy()) 429 | min_feature_size, min_feature_label, labeled_array = find_minimum_feature_size(island_deleted) 430 | 431 | fom_island, _ = opt([island_deleted.flatten()]) 432 | print('min feature size:', min_feature_size) 433 | print('min feature label:', min_feature_label) 434 | print('labeled array shape:', labeled_array.shape) 435 | print('island_deleted shape:', island_deleted.shape) 436 | 437 | wandb.log({ 438 | 'fom (final, island deleted1)': fom_island, 439 | "generated_final, island deleted1": [wandb.Image(1 - island_deleted)], 440 | }) 441 | 442 | # 2) Convert array to further process 443 | array_converted_processed = delete_islands_with_size_1(1 - island_deleted.copy()) 444 | array_original = 1 - array_converted_processed 445 | 446 | min_feature_size2, min_feature_label2, labeled_array2 = find_minimum_feature_size(array_converted_processed) 447 | plt2 = highlight_minimum_island(array_original, min_feature_label2, labeled_array2) 448 | 449 | euler_number_original = euler_number(array_original, connectivity=1) 450 | _, num_islands1 = label_(island_deleted, connectivity=1, return_num=True) 451 | num_holes1 = num_islands1 - euler_number_original 452 | 453 | euler_number_converted = euler_number(array_converted_processed, connectivity=1) 454 | _, num_islands2 = label_(array_converted_processed, connectivity=1, return_num=True) 455 | num_holes2 = num_islands2 - euler_number_converted 456 | 457 | fom_island2, _ = opt([array_original.flatten()]) 458 | 459 | wandb.log({"highlighted, island deleted2": [wandb.Image(plt2)]}) 460 | 461 | plt3 = highlight_minimum_island(array_original, min_feature_label, labeled_array) 462 | 463 | wandb.log({ 464 | 'fom (final, island deleted2)': fom_island2, 465 | "generated_final, island deleted2": [wandb.Image(1 - array_original)], 466 | "highlighted, final": [wandb.Image(plt3)] 467 | }) 468 | 469 | print("array 0,0: ", array_original[0, 0]) 470 | 471 | wandb.log({ 472 | 'mfs (original)': min_feature_size, 473 | 'mfs (converted)': min_feature_size2, 474 | 'number of islands (original)': num_islands1, 475 | 'number of islands (converted)': num_islands2, 476 | 'number of holes (original)': num_holes1, 477 | 'number of holes (converted)': num_holes2, 478 | 'euler_number with 1 connectivity (original)': euler_number_original, 479 | 'euler_number with 1 connectivity (converted)': euler_number_converted, 480 | }) 481 | -------------------------------------------------------------------------------- /guided_diffusion/simulation.py: -------------------------------------------------------------------------------- 1 | import matplotlib.pyplot as plt 2 | 3 | import pickle 4 | import os 5 | # Define global lists 6 | 7 | 8 | plt.rcParams["font.family"] = "sans-serif" 9 | plt.rcParams["font.sans-serif"] = ["DejaVu Sans"] 10 | plt.rcParams["figure.figsize"] = (3.5,3.5) 11 | 12 | params = { 13 | 'axes.labelsize':12, # label 폰트 크기 14 | 'axes.titlesize':12, # 타이틀 폰트 크기 15 | 'xtick.labelsize':10, # x 축 tick label 폰트 크기 16 | 'ytick.labelsize':10, # y 축 tick label 폰트 크기 17 | 'xtick.direction': 'in', # 눈금 표시 방향 (in, out, inout) 18 | 'ytick.direction': 'in', # 눈금 표시 방향 (in, out, inout) 19 | 'lines.markersize': 3, # 마커 사이즈 20 | 'axes.titlepad': 6, # 타이틀과 그래프 사이의 간격 21 | 'axes.labelpad': 4, # 축 label과 그래프 사이의 간격 22 | 'font.size': 12, # font 크기 23 | #'font.sans-serif': 'Arial', # font 설정 24 | 'figure.dpi': 300, # 해상도, vector그래픽의 경우 dpi에 상관없이 깔끔하게 출력됨 25 | 'figure.autolayout': True, # 레이아웃 자동 설정 (그래프의 모든 요소가 figure 내부에 들어가도록 설정) 26 | 'xtick.top': True, # 그래프 위쪽 x축 눈금 표시 27 | 'ytick.right': True, # 그래프 오른쪽 y축 눈금 표시 28 | 'xtick.major.size': 2, # x축 눈금의 길이 29 | 'ytick.major.size': 2, # y축 눈금의 길이 30 | } 31 | 32 | # Save the lists to a file 33 | 34 | def save_lists(new_red_list, new_blue_list, new_green_list): 35 | # Load existing data if file exists 36 | if os.path.exists('lists.pkl'): 37 | with open('lists.pkl', 'rb') as f: 38 | red_list, blue_list, green_list = pickle.load(f) 39 | else: 40 | # If file does not exist, initialize empty lists 41 | red_list, blue_list, green_list = [], [], [] 42 | 43 | # Append new data to existing lists 44 | red_list += new_red_list 45 | blue_list += new_blue_list 46 | green_list += new_green_list 47 | 48 | # Save the updated lists back to the pickle file 49 | with open('lists.pkl', 'wb') as f: 50 | pickle.dump((red_list, blue_list, green_list), f) 51 | 52 | 53 | def CIS_sim(struct_np, t, exp_name, prop_dir ='', 54 | save_inter=False, interval=10, flag_last=False): 55 | import numpy as np 56 | import meep as mp 57 | import meep.adjoint as mpa 58 | import autograd.numpy as npa 59 | import numpy as npa 60 | import os 61 | import matplotlib.pyplot as plt 62 | import wandb 63 | struct_np [struct_np > 1] = 1 64 | struct_np [struct_np < 0 ] = 0 65 | 66 | # print(struct_np) 67 | red_list = [] 68 | blue_list = [] 69 | green_list = [] 70 | 71 | mp.verbosity(0) 72 | Air = mp.Medium(index=1.0) 73 | SiN = mp.Medium(epsilon=4) 74 | SiO2 = mp.Medium(epsilon=2.1) 75 | SiPD = mp.Medium(epsilon=5) 76 | 77 | 78 | um_scale = 1 79 | resolution = 21 80 | 81 | design_region_width = 3 # 디자인 영역 너비 82 | design_region_height = 3 # 디자인 영역 높이 83 | gapop = 0 #################################################################################################### 84 | air_gap = 0 85 | dti = 0.4 86 | subpixelsize = design_region_width/3 - dti 87 | if gapop == 1: 88 | air_gap = dti/2 89 | PDsize = 2 90 | Lpml = 0.5 # PML 영역 크기 91 | Sourcespace = 2 92 | 93 | Sx = design_region_width 94 | Sy = PDsize + design_region_height + Sourcespace + Lpml 95 | cell_size = mp.Vector3(Sx, Sy) 96 | 97 | pml_layers = [mp.PML(thickness = Lpml, direction = mp.Y)] 98 | 99 | # 파장, 주파수 설정 100 | wavelengths = np.linspace(0.40*um_scale, 0.70*um_scale, 31) 101 | frequencies = 1/wavelengths 102 | nf = len(frequencies) # number of frequencies 103 | 104 | 105 | #source 106 | width = 0.4 107 | 108 | fcen_red = 1/(0.65*um_scale) 109 | fwidth_red = fcen_red * width 110 | 111 | fcen_green = 1/(0.55*um_scale) 112 | fwidth_green = fcen_green * width 113 | 114 | fcen_blue = 1/(0.45*um_scale) 115 | fwidth_blue = fcen_blue * width 116 | 117 | src_0 = mp.GaussianSource(frequency=fcen_red, fwidth=fwidth_red, is_integrated=True) 118 | 119 | src_1 = mp.GaussianSource(frequency=fcen_green, fwidth=fwidth_green, is_integrated=True) 120 | 121 | src_2 = mp.GaussianSource(frequency=fcen_blue, fwidth=fwidth_blue, is_integrated=True) 122 | 123 | source_center = mp.Vector3(0, Sy/ 2 - Lpml - Sourcespace / 2, 0) # Source 위치 124 | source_size = mp.Vector3(Sx, 0, 0) 125 | 126 | source = [mp.Source(src_0, component=mp.Ez, size=source_size, center=source_center,), 127 | mp.Source(src_1, component=mp.Ez, size=source_size, center=source_center,), 128 | mp.Source(src_2, component=mp.Ez, size=source_size, center=source_center,),] 129 | 130 | Nx = 64#int(round(design_region_resolution * design_region_width)) + 1 131 | Ny = 64#int(round(design_region_resolution * design_region_height)) + 1 132 | 133 | # 설계 영역과 물질을 바탕으로 설계 영역 설정 134 | design_variables = mp.MaterialGrid(mp.Vector3(Nx, Ny), SiO2, SiN, grid_type="U_MEAN") 135 | design_region = mpa.DesignRegion( 136 | design_variables, 137 | volume=mp.Volume( 138 | center=mp.Vector3(0, - Sy /2 + PDsize + design_region_height / 2, 0), 139 | size=mp.Vector3(design_region_width, design_region_height, 0), 140 | ), 141 | ) 142 | 143 | """ 144 | 145 | """ 146 | # design region과 동일한 size의 Block 생성 147 | geometry = [ 148 | mp.Block( 149 | center=design_region.center, size=design_region.size, material=design_variables 150 | ), 151 | mp.Block( 152 | center=mp.Vector3(0, -Sy/2 + PDsize/2, 0), size=mp.Vector3(Sx, PDsize, 0), material=SiO2 153 | ), 154 | # DTI가 있을 경우 사용 155 | mp.Block( 156 | center=mp.Vector3(-design_region_width/3, -Sy/2 + PDsize/2, 0), size=mp.Vector3(subpixelsize, PDsize, 0), material=SiPD 157 | ), 158 | mp.Block( 159 | center=mp.Vector3(0, -Sy/2 + PDsize/2, 0), size=mp.Vector3(subpixelsize, PDsize, 0), material=SiPD 160 | ), 161 | mp.Block( 162 | center=mp.Vector3(design_region_width/3, -Sy/2 + PDsize/2, 0), size=mp.Vector3(subpixelsize, PDsize, 0), material=SiPD 163 | ) 164 | ] 165 | 166 | # Meep simulation 세팅 167 | sim = mp.Simulation( 168 | cell_size=cell_size, 169 | boundary_layers=pml_layers, 170 | geometry=geometry, 171 | sources=source, 172 | default_material=Air, # 빈공간 173 | resolution=resolution, 174 | k_point = mp.Vector3(0,0,0) # bloch boundary 175 | ) 176 | 177 | 178 | # 모니터 위치와 크기 설정 (focal point) 179 | #monitor_position_0, monitor_size_0 = mp.Vector3(-design_region_width/3, -Sy/2 + PDsize - 0.5/resolution), mp.Vector3(0.01,0) 180 | #monitor_position_1, monitor_size_1 = mp.Vector3(0, -Sy/2 + PDsize - 0.5/resolution), mp.Vector3(0.01,0) 181 | #monitor_position_2, monitor_size_2 = mp.Vector3(design_region_width/3, -Sy/2 + PDsize - 0.5/resolution), mp.Vector3(0.01,0) 182 | #monitor_position_3, monitor_size_3 = mp.Vector3(0, -Sy/2 + PDsize + design_region_height + 0.5/resolution), mp.Vector3(design_region_width,0) 183 | #모니터 위치와 크기 설정 (focal point) 184 | monitor_position_0, monitor_size_0 = mp.Vector3(-design_region_width/3, -Sy/2 + PDsize - 0.5/resolution), mp.Vector3(subpixelsize,0) 185 | monitor_position_1, monitor_size_1 = mp.Vector3(0, -Sy/2 + PDsize - 0.5/resolution), mp.Vector3(subpixelsize,0) 186 | monitor_position_2, monitor_size_2 = mp.Vector3(design_region_width/3, -Sy/2 + PDsize - 0.5/resolution), mp.Vector3(subpixelsize,0) 187 | 188 | 189 | 190 | # FourierFields를 통해 monitor_position에서 monitor_size만큼의 영역에 대한 Fourier transform을 구함 191 | FourierFields_0 = mpa.FourierFields(sim,mp.Volume(center=monitor_position_0,size=monitor_size_0),mp.Ez,yee_grid=True) 192 | 193 | FourierFields_1 = mpa.FourierFields(sim,mp.Volume(center=monitor_position_1,size=monitor_size_1),mp.Ez,yee_grid=True) 194 | 195 | FourierFields_2= mpa.FourierFields(sim,mp.Volume(center=monitor_position_2,size=monitor_size_2),mp.Ez,yee_grid=True) 196 | 197 | 198 | ob_list = [FourierFields_0, FourierFields_1, FourierFields_2,] 199 | 200 | 201 | def J_0(fields_0, fields_1, fields_2): 202 | red = npa.sum(npa.abs(fields_0[21:30,:]) **2) 203 | green = npa.sum(npa.abs(fields_1[11:20,:]) ** 2) 204 | blue = npa.sum(npa.abs(fields_2[1:10,:]) ** 2) 205 | 206 | if isinstance(red, np.floating): 207 | red_ = red 208 | else: 209 | red_ = red._value 210 | 211 | if isinstance(green, np.floating): 212 | green_ = green 213 | else: 214 | green_ = green._value 215 | 216 | if isinstance(blue, np.floating): 217 | blue_ = blue 218 | else: 219 | blue_ = blue._value 220 | 221 | red_list.append(red_) 222 | blue_list.append(blue_) 223 | green_list.append(green_) 224 | save_lists(red_list, blue_list, green_list) 225 | 226 | return blue + green + red 227 | 228 | 229 | # optimization 설정å 230 | opt = mpa.OptimizationProblem( 231 | simulation=sim, 232 | objective_functions=[J_0], 233 | objective_arguments=ob_list, 234 | design_regions=[design_region], 235 | frequencies=frequencies, 236 | decay_by=1e-3, # 모니터에 남아있는 필드 값의 비율 237 | ) 238 | 239 | #struct_random = np.random.choice([0, 1], size=(Nx*Ny)) 240 | 241 | flattened_array = struct_np.flatten() 242 | 243 | # error handling 244 | #try: 245 | opt.update_design([flattened_array]) 246 | #opt.plot2D(True) 247 | #except: 248 | # exit(1) 249 | 250 | 251 | 252 | 253 | #opt.plot2D(fields=mp.Ez) 254 | #plt.savefig('structure.png') 255 | fom, dJ_du = opt([flattened_array]) 256 | g = np.sum(dJ_du, axis=1) 257 | #f0 = f0[0] 258 | #print(f0) 259 | #print(dJ_du) 260 | 261 | return fom, g 262 | 263 | 264 | 265 | 266 | def waveguide_sim(struct_np, t, exp_name, prop_dir='top', 267 | save_inter=False, interval=1, flag_last=False): 268 | import numpy as np 269 | import meep as mp 270 | import meep.adjoint as mpa 271 | import autograd.numpy as npa 272 | import os 273 | import matplotlib.pyplot as plt 274 | import wandb 275 | 276 | assert prop_dir in ['top', 'bottom', 'front'] 277 | mp.verbosity(0) 278 | struct_np [struct_np > 1] = 1 279 | struct_np [struct_np < 0 ] = 0 280 | struct_np = np.squeeze(struct_np) 281 | # print(struct_np) 282 | 283 | #mp.verbosity(0) 284 | Si = mp.Medium(index=3.4) 285 | SiO2 = mp.Medium(index=1.44) 286 | 287 | resolution = 21 288 | 289 | Sx = 10 290 | Sy = 10 291 | cell_size = mp.Vector3(Sx, Sy) 292 | 293 | # pml_layers = [mp.PML(1.0)] 294 | pml_layers = [mp.PML(2.0)] 295 | 296 | fcen = 1 / 1.55 297 | width = 0.2 298 | fwidth = width * fcen 299 | # source_center = [-3.1, 0, 0] 300 | # source_size = mp.Vector3(0, 1, 0) 301 | # source_center = [-3.2, 0, 0] 302 | source_center = [-2.7, 0, 0] 303 | source_size = mp.Vector3(0, 2, 0) 304 | kpoint = mp.Vector3(1, 0, 0) 305 | src = mp.GaussianSource(frequency=fcen, fwidth=fwidth) 306 | source = [ 307 | mp.EigenModeSource( 308 | src, 309 | eig_band=1, 310 | direction=mp.NO_DIRECTION, 311 | eig_kpoint=kpoint, 312 | size=source_size, 313 | center=source_center, 314 | ) 315 | ] 316 | 317 | design_region_resolution = 21 318 | Nx = 64#design_region_resolution +1 319 | Ny = 64#design_region_resolution +1 320 | 321 | design_variables = mp.MaterialGrid(mp.Vector3(Nx, Ny), SiO2, Si, grid_type="U_MEAN") 322 | design_region = mpa.DesignRegion( 323 | design_variables, volume=mp.Volume(center=mp.Vector3(), size=mp.Vector3(3, 3, 0)) 324 | ) 325 | 326 | """ 327 | 328 | """ 329 | if prop_dir == 'top': 330 | geometry = [ 331 | mp.Block( 332 | center=mp.Vector3(x=-Sx / 4), material=Si, size=mp.Vector3(Sx / 2, 1, 0) 333 | ), # horizontal waveguide: left (origin) 334 | mp.Block( 335 | center=mp.Vector3(y=Sy / 4), material=Si, size=mp.Vector3(1, Sy / 2, 0) 336 | ), # vertical waveguide: top 337 | mp.Block( 338 | center=design_region.center, size=design_region.size, material=design_variables 339 | ), # design region 340 | # The commented lines above impose symmetry by overlapping design region with the same design variable. However, 341 | # currently there is an issue of doing that; We give an alternative approach to impose symmetry in later tutorials. 342 | # See https://github.com/NanoComp/meep/issues/1984 and https://github.com/NanoComp/meep/issues/2093 343 | ] 344 | 345 | elif prop_dir == 'bottom': 346 | geometry = [ 347 | mp.Block( 348 | center=mp.Vector3(x=-Sx / 4), material=Si, size=mp.Vector3(Sx / 2, 1, 0) 349 | ), # horizontal waveguide: left (origin) 350 | mp.Block( 351 | center=mp.Vector3(y=-Sy / 4), material=Si, size=mp.Vector3(1, Sy / 2, 0) 352 | ), # vertical waveguide: bottom 353 | mp.Block( 354 | center=design_region.center, size=design_region.size, material=design_variables 355 | ), # design region 356 | # The commented lines above impose symmetry by overlapping design region with the same design variable. However, 357 | # currently there is an issue of doing that; We give an alternative approach to impose symmetry in later tutorials. 358 | # See https://github.com/NanoComp/meep/issues/1984 and https://github.com/NanoComp/meep/issues/2093 359 | ] 360 | elif prop_dir == 'front': 361 | geometry = [ 362 | mp.Block( 363 | center=mp.Vector3(x=-Sx / 4), material=Si, size=mp.Vector3(Sx / 2, 1, 0) 364 | ), # horizontal waveguide: left (origin) 365 | mp.Block( 366 | center=mp.Vector3(x=Sx / 4), material=Si, size=mp.Vector3(Sx / 2, 1, 0) 367 | ), # horizontal waveguide: front 368 | mp.Block( 369 | center=design_region.center, size=design_region.size, material=design_variables 370 | ), # design region 371 | # The commented lines above impose symmetry by overlapping design region with the same design variable. However, 372 | # currently there is an issue of doing that; We give an alternative approach to impose symmetry in later tutorials. 373 | # See https://github.com/NanoComp/meep/issues/1984 and https://github.com/NanoComp/meep/issues/2093 374 | ] 375 | 376 | 377 | 378 | 379 | sim = mp.Simulation( 380 | cell_size=cell_size, 381 | boundary_layers=pml_layers, 382 | geometry=geometry, 383 | sources=source, 384 | eps_averaging=True, 385 | subpixel_tol=1e-4, 386 | resolution=resolution, 387 | ) 388 | 389 | 390 | TE_front = mpa.EigenmodeCoefficient( 391 | sim, mp.Volume(center=mp.Vector3(2.5, 0, 0), size=mp.Vector3(y=2)), mode=1 392 | ) 393 | TE_top = mpa.EigenmodeCoefficient( 394 | sim, mp.Volume(center=mp.Vector3(0, 2.5, 0), size=mp.Vector3(x=2)), mode=1 395 | ) 396 | TE_bottom = mpa.EigenmodeCoefficient( 397 | sim, mp.Volume(center=mp.Vector3(0, -2.5, 0), size=mp.Vector3(x=2)), mode=1, forward=False 398 | ) 399 | TE_o = mpa.EigenmodeCoefficient( 400 | sim, mp.Volume(center=mp.Vector3(-2.5, 0, 0), size=mp.Vector3(y=2)), mode=1 401 | ) 402 | 403 | ob_list = [TE_o] 404 | 405 | if prop_dir == 'top': 406 | ob_list.append(TE_top) 407 | # ob_list = [TE_o, TE_top] 408 | elif prop_dir == 'bottom': 409 | ob_list.append(TE_bottom) 410 | # ob_list = [TE_o, TE_bottom] 411 | elif prop_dir == 'front': 412 | ob_list.append(TE_front) 413 | # ob_list = [TE_o, TE_front] 414 | 415 | def J(source, target): 416 | return npa.abs(target / source) ** 2 417 | 418 | opt = mpa.OptimizationProblem( 419 | simulation=sim, 420 | objective_functions=J, 421 | objective_arguments=ob_list, 422 | design_regions=[design_region], 423 | fcen=fcen, 424 | df=0, 425 | nf=1, 426 | ) 427 | 428 | #struct_random = np.random.choice([0, 1], size=(Nx*Ny)) 429 | 430 | flattened_array = struct_np.flatten() 431 | opt.update_design([flattened_array]) 432 | #opt.plot2D(True) 433 | 434 | #if save_inter and (t % interval == 0): 435 | # os.makedirs(f'figures/{exp_name}', exist_ok=True) 436 | # plt.savefig(f'figures/{exp_name}/structure_t={wandb.config.tsr-t}.png') 437 | 438 | #if flag_last: 439 | # os.makedirs(f'figures/{exp_name}', exist_ok=True) 440 | # plt.savefig(f'figures/{exp_name}/structure_final.png') 441 | 442 | 443 | # if filename is not None: 444 | # plt.savefig(f'structure_{filename}.png') 445 | # else: 446 | # plt.savefig('sample/structure.png') 447 | #opt.plot2D(fields=mp.Ez) 448 | #plt.savefig('structure.png') 449 | 450 | fom, g = opt([flattened_array]) 451 | fom = fom[0] 452 | #f0 = f0[0] 453 | #print(f0) 454 | #print(dJ_du) 455 | return fom, g -------------------------------------------------------------------------------- /baseline_algorithms/[adjoint_optimization] GA_CIS.py: -------------------------------------------------------------------------------- 1 | # Design CIS 2D color router 2 | # Resolution 25 3 | 4 | # ## 1. Simulation Environment 5 | import copy 6 | import meep as mp 7 | import meep.adjoint as mpa 8 | import numpy as np 9 | import nlopt 10 | from autograd import numpy as npa 11 | from autograd import tensor_jacobian_product, grad 12 | from matplotlib import pyplot as plt 13 | from scipy.ndimage import label, binary_dilation 14 | import os 15 | from skimage.measure import euler_number 16 | from skimage.measure import label as label_ 17 | 18 | 19 | import wandb 20 | wandb.init(project="colorrouter_GD") 21 | mp.verbosity(1) 22 | 23 | 24 | design_dir = "./CIS_design_GD/" 25 | 26 | # Create directory if it doesn't exist 27 | if not os.path.exists(design_dir): 28 | os.makedirs(design_dir) 29 | 30 | plt.rcParams["font.family"] = "sans-serif" 31 | plt.rcParams["font.sans-serif"] = ["DejaVu Sans"] 32 | 33 | plt.rcParams["figure.figsize"] = (3.5,3.5) 34 | 35 | 36 | 37 | params = { 38 | 'axes.labelsize':9, # label font size 39 | 'axes.titlesize':9, # title font size 40 | 'xtick.labelsize':9, # x axis tick label font size 41 | 'ytick.labelsize':9, # y axis tick label font size 42 | 'xtick.direction': 'in', # tick direction (in, out, inout) 43 | 'ytick.direction': 'in', # tick direction (in, out, inout) 44 | 'lines.markersize': 3, # marker size 45 | 'axes.titlepad': 6, # spacing between title and graph 46 | 'axes.labelpad': 4, # spacing between axis label and graph 47 | 'font.size': 9, # font size 48 | 'figure.dpi': 300, # resolution, vector graphics output cleanly regardless of dpi 49 | 'figure.autolayout': True, # automatic layout (all graph elements fit inside figure) 50 | 'xtick.top': True, # show x-axis ticks on top 51 | 'ytick.right': True, # show y-axis ticks on right 52 | 'xtick.major.size': 2, # x-axis tick length 53 | 'ytick.major.size': 2, # y-axis tick length 54 | } 55 | 56 | 57 | 58 | # Function to delete every island with size 1 59 | def delete_islands_with_size_1(array): 60 | # Label the connected components 61 | labeled_array, num_features = label(array) 62 | 63 | # Identify the features (islands) with size 1 64 | islands_to_delete = [i for i in range(1, num_features + 1) if np.sum(labeled_array == i) == 1] 65 | 66 | # Delete these islands 67 | for island in islands_to_delete: 68 | array[labeled_array == island] = 0 69 | 70 | if array.ndim > 2: 71 | array = np.squeeze(array, axis=2) 72 | else: 73 | pass 74 | return array 75 | 76 | 77 | # Function to find the minimum feature size and its location 78 | def find_minimum_feature_size(array): 79 | # Label the connected components 80 | labeled_array, num_features = label(array) 81 | 82 | # Get the sizes of all features 83 | feature_sizes = [(i, np.sum(labeled_array == i)) for i in range(1, num_features + 1)] 84 | 85 | # Find the minimum feature size and its label 86 | min_feature = min(feature_sizes, key=lambda x: x[1]) if feature_sizes else (0, 0) 87 | 88 | return min_feature[1], min_feature[0], labeled_array 89 | 90 | 91 | 92 | def highlight_minimum_island(array, min_label, labeled_array, title): 93 | # Find the coordinates of the minimum island 94 | coords = np.argwhere(labeled_array == min_label) 95 | if coords.size == 0: 96 | return array 97 | 98 | # Create an empty mask for the boundary 99 | boundary_mask = np.zeros_like(array) 100 | 101 | # Fill the mask with the minimum island 102 | boundary_mask[labeled_array == min_label] = 1 103 | 104 | # Dilate the mask and subtract the original to get the boundary 105 | dilated_mask = binary_dilation(boundary_mask) 106 | boundary = dilated_mask - boundary_mask 107 | 108 | # Create a plot 109 | fig, ax = plt.subplots() 110 | ax.imshow(1-array, cmap='gray') # black: 1, white: 0 111 | ax.imshow(boundary, cmap='Reds', alpha=0.5) # Overlay the red boundary with higher alpha 112 | ax.set_title(title) 113 | 114 | # Adjust boundary line thickness and make it more vivid 115 | boundary_indices = np.argwhere(boundary) 116 | for y, x in boundary_indices: 117 | rect = plt.Rectangle((x - 0.5, y - 0.5), 1, 1, edgecolor='red', facecolor='red', linewidth=1) 118 | ax.add_patch(rect) 119 | 120 | return plt 121 | 122 | 123 | plt.rcParams.update(params) 124 | 125 | # scaling & refractive index 126 | um_scale = 1 127 | 128 | Air = mp.Medium(index=1.0) 129 | SiN = mp.Medium(epsilon=4) 130 | SiO2 = mp.Medium(epsilon=2.1) 131 | TiO2 = mp.Medium(epsilon=7) 132 | SiPD = mp.Medium(epsilon=5) 133 | 134 | # Design space 135 | design_region_width = 3 # design region width 136 | design_region_height = 3#4 # design region height 137 | 138 | # Resolution and size settings 139 | resolution = 21 140 | gapop = 0 141 | air_gap = 0 142 | dti = 0.4 143 | subpixelsize = design_region_width/3 - dti 144 | if gapop == 1: 145 | air_gap = dti/2 146 | PDsize = 2 147 | Lpml = 0.5 # PML region size 148 | pml_layers = [mp.PML(thickness = Lpml, direction = mp.Y)] 149 | Sourcespace = 2 150 | 151 | 152 | 153 | # Total space 154 | Sx = design_region_width 155 | Sy = PDsize + design_region_height + Sourcespace + Lpml 156 | cell_size = mp.Vector3(Sx, Sy) 157 | 158 | # Wavelength and frequency settings 159 | wavelengths = np.linspace(0.40*um_scale, 0.70*um_scale, 31) 160 | frequencies = 1/wavelengths 161 | nf = len(frequencies) # number of frequencies 162 | 163 | 164 | # Fabrication Constraints settings 165 | 166 | minimum_length = 0.895 # minimum length scale (microns) 167 | eta_i = 0.5 # blueprint (or intermediate) design field thresholding point (between 0 and 1) 168 | filter_radius = minimum_length 169 | design_region_resolution = int(resolution) 170 | print(filter_radius) 171 | 172 | 173 | # Source settings 174 | 175 | width = 0.4 176 | 177 | fcen_red = 1/(0.65*um_scale) 178 | fwidth_red = fcen_red * width 179 | 180 | fcen_green = 1/(0.55*um_scale) 181 | fwidth_green = fcen_green * width 182 | 183 | fcen_blue = 1/(0.45*um_scale) 184 | fwidth_blue = fcen_blue * width 185 | 186 | src_0 = mp.GaussianSource(frequency=fcen_red, fwidth=fwidth_red, is_integrated=True) 187 | 188 | src_1 = mp.GaussianSource(frequency=fcen_green, fwidth=fwidth_green, is_integrated=True) 189 | 190 | src_2 = mp.GaussianSource(frequency=fcen_blue, fwidth=fwidth_blue, is_integrated=True) 191 | 192 | source_center = mp.Vector3(0, Sy/ 2 - Lpml - Sourcespace / 2, 0) # Source position 193 | source_size = mp.Vector3(Sx, 0, 0) 194 | 195 | source = [mp.Source(src_0, component=mp.Ez, size=source_size, center=source_center,), 196 | mp.Source(src_1, component=mp.Ez, size=source_size, center=source_center,), 197 | mp.Source(src_2, component=mp.Ez, size=source_size, center=source_center,),] 198 | 199 | 200 | 201 | 202 | # Design region pixels - determined by resolution and design region 203 | Nx = 64 204 | Ny = 64 205 | print('Nx: ', Nx) 206 | print('Ny: ', Ny) 207 | 208 | wandb.config.update({"minimum length": minimum_length, "resolution": resolution, 209 | "Sx": Sx, "Sy": Sy, "Nx": Nx, "Ny": Ny, 210 | "eta_i": eta_i, 211 | }) 212 | # Design region and material based design region setup 213 | design_variables = mp.MaterialGrid(mp.Vector3(Nx, Ny), SiO2, SiN, grid_type="U_MEAN") 214 | design_region = mpa.DesignRegion( 215 | design_variables, 216 | volume=mp.Volume( 217 | center=mp.Vector3(0, - Sy /2 + PDsize + design_region_height / 2, 0), 218 | size=mp.Vector3(design_region_width-air_gap*2, design_region_height, 0), 219 | ), 220 | ) 221 | 222 | 223 | 224 | # Use conic_filter function and simple_2d_filter function from filter.py 225 | def mapping(x, eta, beta): 226 | # filter 227 | 228 | try: 229 | filtered_field = mpa.conic_filter( 230 | x, 231 | filter_radius, 232 | design_region_width, 233 | design_region_height, 234 | design_region_resolution, 235 | ) 236 | x = x.flatten() 237 | print('filtered_field type: ', type(filtered_field)) 238 | print("filtered_field shape:", getattr(filtered_field, 'shape', 'No shape attribute')) 239 | 240 | if not isinstance(filtered_field, np.ndarray): 241 | raise TypeError("Expected filtered_field to be an ndarray, got {}".format(type(filtered_field))) 242 | 243 | beta = float(beta) 244 | eta = float(eta) 245 | 246 | projected_field = mpa.tanh_projection(filtered_field, beta, eta) 247 | 248 | except Exception as e: 249 | print("Error in mapping function:", e) 250 | raise 251 | # projection 252 | # limit output to 0 ~ 1 253 | 254 | 255 | # interpolate to actual materials 256 | return projected_field.flatten() 257 | 258 | 259 | # Create Block with same size as design region 260 | geometry = [ 261 | mp.Block( 262 | center=design_region.center, size=design_region.size, material=design_variables 263 | ), 264 | mp.Block( 265 | center=mp.Vector3(0, -Sy/2 + PDsize/2, 0), size=mp.Vector3(Sx, PDsize, 0), material=SiO2 266 | ), 267 | # Use when DTI is present 268 | mp.Block( 269 | center=mp.Vector3(-design_region_width/3, -Sy/2 + PDsize/2, 0), size=mp.Vector3(subpixelsize, PDsize, 0), material=SiPD 270 | ), 271 | mp.Block( 272 | center=mp.Vector3(0, -Sy/2 + PDsize/2, 0), size=mp.Vector3(subpixelsize, PDsize, 0), material=SiPD 273 | ), 274 | mp.Block( 275 | center=mp.Vector3(design_region_width/3, -Sy/2 + PDsize/2, 0), size=mp.Vector3(subpixelsize, PDsize, 0), material=SiPD 276 | ) 277 | ] 278 | 279 | 280 | 281 | # Meep simulation setup 282 | sim = mp.Simulation( 283 | cell_size=cell_size, 284 | boundary_layers=pml_layers, 285 | geometry=geometry, 286 | sources=source, 287 | default_material=Air, # empty space 288 | resolution=resolution, 289 | k_point = mp.Vector3(0,0,0) # bloch boundary 290 | ) 291 | 292 | ############################################################################################################################### 293 | # ## 2. Optimization Environment 294 | 295 | 296 | # Monitor position and size settings (focal point) 297 | monitor_position_0, monitor_size_0 = mp.Vector3(-design_region_width/3, -Sy/2 + PDsize - 0.5/resolution), mp.Vector3(subpixelsize,0) 298 | monitor_position_1, monitor_size_1 = mp.Vector3(0, -Sy/2 + PDsize - 0.5/resolution), mp.Vector3(subpixelsize,0) 299 | monitor_position_2, monitor_size_2 = mp.Vector3(design_region_width/3, -Sy/2 + PDsize - 0.5/resolution), mp.Vector3(subpixelsize,0) 300 | 301 | 302 | 303 | # Get Fourier transform for region of monitor_size at monitor_position using FourierFields 304 | FourierFields_0 = mpa.FourierFields(sim,mp.Volume(center=monitor_position_0,size=monitor_size_0),mp.Ez,yee_grid=True) 305 | 306 | FourierFields_1 = mpa.FourierFields(sim,mp.Volume(center=monitor_position_1,size=monitor_size_1),mp.Ez,yee_grid=True) 307 | 308 | FourierFields_2= mpa.FourierFields(sim,mp.Volume(center=monitor_position_2,size=monitor_size_2),mp.Ez,yee_grid=True) 309 | 310 | 311 | ob_list = [FourierFields_0, FourierFields_1, FourierFields_2,] 312 | 313 | 314 | 315 | 316 | fred = [] 317 | fgreen = [] 318 | fblue = [] 319 | # J : Objective function 320 | # Take the squared absolute value of the Ez component at the center of the monitor measured by FourierFields 321 | # [frequency index, monitor index] 322 | 323 | 324 | flag = 0 325 | def J_0(fields_0, fields_1, fields_2): 326 | red = npa.sum(npa.abs(fields_0[21:30,:]) **2) 327 | green = npa.sum(npa.abs(fields_1[11:20,:]) ** 2) 328 | blue = npa.sum(npa.abs(fields_2[1:10,:]) ** 2) 329 | 330 | 331 | 332 | redfactor = 1 333 | greenfactor = 1 334 | bluefactor = 1 335 | 336 | if isinstance(red, np.floating): 337 | red_ = red 338 | else: 339 | red_ = red._value 340 | 341 | if isinstance(green, np.floating): 342 | green_ = green 343 | else: 344 | green_ = green._value 345 | 346 | if isinstance(blue, np.floating): 347 | blue_ = blue 348 | else: 349 | blue_ = blue._value 350 | 351 | 352 | 353 | fred.append(red_/redfactor) 354 | fgreen.append(green_/greenfactor) 355 | fblue.append(blue_/bluefactor) 356 | 357 | return blue/bluefactor + green/greenfactor + red/redfactor 358 | 359 | 360 | 361 | 362 | # Optimization settings 363 | opt = mpa.OptimizationProblem( 364 | simulation=sim, 365 | objective_functions=[J_0], 366 | objective_arguments=ob_list, 367 | design_regions=[design_region], 368 | frequencies=frequencies, 369 | decay_by=1e-1, # ratio of field values remaining at monitor 370 | ) 371 | 372 | 373 | evaluation_history = [] 374 | cur_iter = [0] 375 | numevl = 1 376 | 377 | n = Nx * Ny # number of parameters 378 | 379 | 380 | 381 | 382 | # lower and upper bounds (upper: 1, lower: 0) 383 | lb = np.zeros((Nx * Ny,)) 384 | ub = np.ones((Nx * Ny,)) 385 | LR = 0.01 386 | 387 | 388 | cur_beta = 2 389 | beta_scale = 2 390 | num_betas = 140 391 | wandb.config.update({ 392 | 'beta_scale': beta_scale, 393 | 'num_betas': num_betas}) 394 | 395 | 396 | x = np.ones(n)*0.5 397 | 398 | wandb.run.name = "min"+str(minimum_length)+"_"+"eta"+str(eta_i)+"_"+"num_betas"+str(num_betas) 399 | 400 | 401 | for iters in range(num_betas): 402 | x [x>1] = 1 403 | x [x<0] = 0 404 | 405 | f0, dJ_du = opt([mapping(x, eta_i, cur_beta)]) 406 | adjoint_gradient = np.sum(dJ_du, axis =1 ) 407 | reshaped_gradients = adjoint_gradient.reshape(adjoint_gradient.shape[0], -1) 408 | 409 | # Computing the norm of each reshaped row 410 | norms = np.linalg.norm(reshaped_gradients, axis=1) 411 | 412 | # Calculating the mean of the norms 413 | adjgrad_norm = norms.mean() 414 | 415 | 416 | print(adjoint_gradient) 417 | print(type(adjoint_gradient)) 418 | learning_rate = LR/adjgrad_norm 419 | 420 | 421 | x_image = x.reshape(Nx, Ny,1) 422 | print(x_image.shape) 423 | 424 | wandb.log({ 425 | 'fom ': np.real(f0), 426 | 'adjgrad norm': adjgrad_norm, 427 | 'learning rate': learning_rate, 428 | "generated": [wandb.Image(x_image, caption=f"Iteration {iters}")] 429 | 430 | }), 431 | 432 | x = x + learning_rate * adjoint_gradient 433 | cur_beta = cur_beta * beta_scale 434 | 435 | ## final 436 | 437 | x [x>0.5] = 1 438 | x [x==0.5] = 1 439 | x [x<0.5] = 0 440 | 441 | x_final = mapping(x, eta_i, mp.inf) 442 | 443 | f0, dJ_du = opt([x_final]) 444 | adjoint_gradient = np.sum(dJ_du, axis =1 ) 445 | 446 | x_image = x.reshape(Nx, Ny,1) 447 | 448 | wandb.log({ 449 | 'final_fom ': np.real(f0), 450 | "generated_final": [wandb.Image(1-x_image, caption=f"generated_final")] 451 | 452 | }), 453 | 454 | sensitivity = [0] 455 | 456 | sensitivity[0] = adjoint_gradient 457 | 458 | fig = plt.figure(figsize=(20,20)) 459 | plt.imshow(np.squeeze(np.abs(sensitivity[0].reshape(Nx, Ny)))) 460 | plt.xlabel("x") 461 | plt.ylabel("y") 462 | cmap = plt.cm.get_cmap() 463 | colormapping = plt.cm.ScalarMappable(cmap=cmap) 464 | cbar = fig.colorbar(colormapping, ax=plt.gca()) 465 | 466 | wandb.log({"sensitivity": plt}) 467 | 468 | 469 | x_final = x_final.reshape(Nx, Ny) 470 | 471 | island_deleted = delete_islands_with_size_1(x_final.copy()) 472 | min_feature_size, min_feature_label, labeled_array = find_minimum_feature_size(island_deleted) 473 | 474 | fom_island, _ = opt([ 475 | island_deleted.flatten()]) 476 | print('min feature size: ', min_feature_size) 477 | print('min feature label: ', min_feature_label) 478 | print('labeled array: ', labeled_array) 479 | print('island_deleted: ', island_deleted.shape) 480 | 481 | wandb.log({ 482 | 'fom (final, island deleted1)': fom_island, 483 | "generated_final, island deleted1": [wandb.Image(1-island_deleted, caption=f"generated_final_island_deleted")], 484 | }) 485 | array_converted_processed = delete_islands_with_size_1(1-island_deleted.copy()) 486 | array_original = 1- array_converted_processed 487 | 488 | min_feature_size2, min_feature_label2, labeled_array2 = find_minimum_feature_size(array_converted_processed) 489 | plt2 = highlight_minimum_island(array_original, min_feature_label2, labeled_array2, title='Minimum Feature 0s') 490 | euler_number_original = euler_number(array_original, connectivity=1) 491 | _, num_islands1 = label_(island_deleted, connectivity=1, return_num=True) 492 | num_holes1 = num_islands1 - euler_number_original 493 | 494 | euler_number_converted = euler_number(array_converted_processed, connectivity=1) 495 | _, num_islands2 = label_(array_converted_processed, connectivity=1, return_num=True) 496 | num_holes2 = num_islands2 - euler_number_converted 497 | 498 | 499 | fom_island2, _ = opt([ 500 | array_original.flatten()] 501 | ) 502 | wandb.log({ 503 | "highlighted, island deleted2": [wandb.Image(plt2, caption=f"generated_final_island_deleted_highlghted")], 504 | 505 | }) 506 | plt3 = highlight_minimum_island(array_original, min_feature_label, labeled_array, title='Minimum Feature 1s') 507 | 508 | 509 | wandb.log({ 510 | 'fom (final, island deleted2)': fom_island2, 511 | "generated_final, island deleted2": [wandb.Image(1-array_original, caption=f"generated_final_island_deleted")], 512 | "highlighted, final": [wandb.Image(plt3, caption=f"generated_final_island_deleted_highlghted")] 513 | }) 514 | 515 | print("array 0,0: ", array_original[0,0]) 516 | 517 | 518 | wandb.log({ 519 | 'mfs (original)': min_feature_size, 520 | 'mfs (converted)': min_feature_size2, 521 | 'number of islands (original)': num_islands1, 522 | 'number of islands (converted)':num_islands2, 523 | 'number of holes (original)' : num_holes1, 524 | 'number of holes (converted)': num_holes2, 525 | 'euler_number with 1 connectivity (original)': euler_number_original, 526 | 'euler_number with 1 connectivity (converted)': euler_number_converted, 527 | }) 528 | 529 | -------------------------------------------------------------------------------- /baseline_algorithms/[adjoint_optimization] nlopt_CIS.py: -------------------------------------------------------------------------------- 1 | # === CIS 2D Color Router Design — Resolution 25 === 2 | 3 | # === 1. Imports & Setup === 4 | import os 5 | import copy 6 | import numpy as np 7 | import nlopt 8 | import meep as mp 9 | import meep.adjoint as mpa 10 | from autograd import numpy as npa 11 | from autograd import tensor_jacobian_product, grad 12 | import torch as th 13 | from scipy.ndimage import label, binary_dilation 14 | from skimage.measure import euler_number, label as sklabel 15 | import wandb 16 | import matplotlib.pyplot as plt 17 | 18 | # Initialize Weights & Biases 19 | wandb.init(project="colorrouter_adjoint") 20 | 21 | # Remove previous log file 22 | if os.path.exists('lists.pkl'): 23 | os.remove('lists.pkl') 24 | 25 | # Matplotlib defaults 26 | plt.rcParams.update({ 27 | 'figure.figsize': (3.5, 3.5), 28 | 'figure.dpi': 300, 29 | 'figure.autolayout': True, 30 | 'font.family': 'sans-serif', 31 | 'font.sans-serif': ['DejaVu Sans'], 32 | 'axes.labelsize': 12, 33 | 'axes.titlesize': 12, 34 | 'axes.titlepad': 6, 35 | 'axes.labelpad': 4, 36 | 'xtick.labelsize': 10, 37 | 'ytick.labelsize': 10, 38 | 'xtick.direction': 'in', 39 | 'ytick.direction': 'in', 40 | 'xtick.top': True, 41 | 'ytick.right': True, 42 | 'xtick.major.size': 2, 43 | 'ytick.major.size': 2, 44 | 'lines.markersize': 3, 45 | }) 46 | 47 | # === 2. Utility Functions === 48 | 49 | # Function to delete every island with size 1 50 | def delete_islands_with_size_1(array): 51 | # Label the connected components 52 | labeled_array, num_features = label(array) 53 | 54 | # Identify the features (islands) with size 1 55 | islands_to_delete = [i for i in range(1, num_features + 1) if np.sum(labeled_array == i) == 1] 56 | 57 | # Delete these islands 58 | for island in islands_to_delete: 59 | array[labeled_array == island] = 0 60 | 61 | if array.ndim > 2: 62 | array = np.squeeze(array, axis=2) 63 | else: 64 | pass 65 | return array 66 | 67 | 68 | # Function to find the minimum feature size and its location 69 | def find_minimum_feature_size(array): 70 | # Label the connected components 71 | labeled_array, num_features = label(array) 72 | 73 | # Get the sizes of all features 74 | feature_sizes = [(i, np.sum(labeled_array == i)) for i in range(1, num_features + 1)] 75 | 76 | # Find the minimum feature size and its label 77 | min_feature = min(feature_sizes, key=lambda x: x[1]) if feature_sizes else (0, 0) 78 | 79 | return min_feature[1], min_feature[0], labeled_array 80 | 81 | 82 | 83 | def highlight_minimum_island(array, min_label, labeled_array): 84 | # Find the coordinates of the minimum island 85 | coords = np.argwhere(labeled_array == min_label) 86 | if coords.size == 0: 87 | return array 88 | 89 | # Create an empty mask for the boundary 90 | boundary_mask = np.zeros_like(array) 91 | 92 | # Fill the mask with the minimum island 93 | boundary_mask[labeled_array == min_label] = 1 94 | 95 | # Dilate the mask and subtract the original to get the boundary 96 | dilated_mask = binary_dilation(boundary_mask) 97 | boundary = dilated_mask - boundary_mask 98 | 99 | # Create a plot 100 | fig, ax = plt.subplots() 101 | ax.imshow(1-array, cmap='gray') # black: 1, white: 0 102 | ax.imshow(boundary, cmap='Reds', alpha=0.5) # Overlay the red boundary with higher alpha 103 | 104 | # Adjust boundary line thickness and make it more vivid 105 | boundary_indices = np.argwhere(boundary) 106 | for y, x in boundary_indices: 107 | rect = plt.Rectangle((x - 0.5, y - 0.5), 1, 1, edgecolor='red', facecolor='red', linewidth=1) 108 | ax.add_patch(rect) 109 | 110 | return plt 111 | 112 | # === 3. Algorithm Selection === 113 | ALGO_MMA = nlopt.LD_MMA 114 | ALGO_SLSQP = nlopt.LD_SLSQP 115 | ALGORITHMS = { 116 | ALGO_MMA: 'MMA', 117 | ALGO_SLSQP: 'SLSQP', 118 | } 119 | 120 | def get_algorithm_name(algorithm) -> str: 121 | """ 122 | Return the human-readable name of an NLopt algorithm. 123 | """ 124 | return ALGORITHMS.get(algorithm, 'Unknown') 125 | 126 | # Select optimization algorithm 127 | algorithm = ALGO_MMA 128 | 129 | 130 | 131 | mp.verbosity(1) 132 | 133 | 134 | design_dir = "./CIS_design/" 135 | 136 | # Create directory if it doesn't exist 137 | if not os.path.exists(design_dir): 138 | os.makedirs(design_dir) 139 | 140 | # scaling & refractive index 141 | um_scale = 1 142 | 143 | mp.verbosity(0) 144 | 145 | Air = mp.Medium(index=1.0) 146 | SiN = mp.Medium(epsilon=4) 147 | SiO2 = mp.Medium(epsilon=2.1) 148 | TiO2 = mp.Medium(epsilon=7) 149 | SiPD = mp.Medium(epsilon=5) 150 | 151 | # Design space 152 | design_region_width = 3 # Design region width 153 | design_region_height = 3#4 # Design region height 154 | 155 | # Resolution and size settings 156 | resolution = 21 157 | gapop = 0 #################################################################################################### 158 | air_gap = 0 159 | dti = 0.4 160 | subpixelsize = design_region_width/3 - dti 161 | if gapop == 1: 162 | air_gap = dti/2 163 | PDsize = 2 164 | Lpml = 0.5 # PML region size 165 | pml_layers = [mp.PML(thickness = Lpml, direction = mp.Y)] 166 | Sourcespace = 2 167 | 168 | 169 | 170 | # Total space 171 | Sx = design_region_width 172 | Sy = PDsize + design_region_height + Sourcespace + Lpml 173 | cell_size = mp.Vector3(Sx, Sy) 174 | 175 | # Wavelength, frequency settings 176 | wavelengths = np.linspace(0.40*um_scale, 0.70*um_scale, 31) 177 | frequencies = 1/wavelengths 178 | nf = len(frequencies) # number of frequencies 179 | 180 | 181 | # Fabrication Constraints settings 182 | 183 | minimum_length = 0.224 # minimum length scale (microns) 184 | eta_i = 0.5 # blueprint (or intermediate) design field thresholding point (between 0 and 1) 185 | eta_e = 0.75 # erosion design field thresholding point (between 0 and 1) 186 | eta_d = 1 - eta_e # dilation design field thresholding point (between 0 and 1) 187 | filter_radius = minimum_length#mpa.get_conic_radius_from_eta_e(minimum_length, eta_e) 188 | design_region_resolution = int(resolution) 189 | print(filter_radius) 190 | 191 | 192 | # Source settings 193 | 194 | width = 0.4 195 | 196 | fcen_red = 1/(0.65*um_scale) 197 | fwidth_red = fcen_red * width 198 | 199 | fcen_green = 1/(0.55*um_scale) 200 | fwidth_green = fcen_green * width 201 | 202 | fcen_blue = 1/(0.45*um_scale) 203 | fwidth_blue = fcen_blue * width 204 | 205 | src_0 = mp.GaussianSource(frequency=fcen_red, fwidth=fwidth_red, is_integrated=True) 206 | 207 | src_1 = mp.GaussianSource(frequency=fcen_green, fwidth=fwidth_green, is_integrated=True) 208 | 209 | src_2 = mp.GaussianSource(frequency=fcen_blue, fwidth=fwidth_blue, is_integrated=True) 210 | 211 | source_center = mp.Vector3(0, Sy/ 2 - Lpml - Sourcespace / 2, 0) # Source position 212 | source_size = mp.Vector3(Sx, 0, 0) 213 | 214 | source = [mp.Source(src_0, component=mp.Ez, size=source_size, center=source_center,), 215 | mp.Source(src_1, component=mp.Ez, size=source_size, center=source_center,), 216 | mp.Source(src_2, component=mp.Ez, size=source_size, center=source_center,),] 217 | 218 | 219 | 220 | 221 | # Design region pixels - determined by resolution and design region 222 | Nx = 64#int(round(design_region_resolution * design_region_width)) + 1 223 | Ny = 64#int(round(design_region_resolution * design_region_height)) + 1 224 | print('Nx: ', Nx) 225 | print('Ny: ', Ny) 226 | 227 | 228 | wandb.config.update({"minimum length": minimum_length, "resolution": resolution, 229 | "Sx": Sx, "Sy": Sy, "Nx": Nx, "Ny": Ny, 230 | "eta_i": eta_i, "eta_e": eta_e, "eta_d": eta_d 231 | }) 232 | # Design region and material-based design region setup 233 | design_variables = mp.MaterialGrid(mp.Vector3(Nx, Ny), SiO2, SiN, grid_type="U_MEAN") 234 | design_region = mpa.DesignRegion( 235 | design_variables, 236 | volume=mp.Volume( 237 | center=mp.Vector3(0, - Sy /2 + PDsize + design_region_height / 2, 0), 238 | size=mp.Vector3(design_region_width-air_gap*2, design_region_height, 0), 239 | ), 240 | ) 241 | 242 | 243 | 244 | # Use conic_filter function and simple_2d_filter function from filter.py 245 | def mapping(x, eta, beta): 246 | # filter 247 | x = x.flatten() 248 | 249 | filtered_field = mpa.conic_filter( 250 | x, 251 | filter_radius, 252 | design_region_width, 253 | design_region_height, 254 | design_region_resolution, 255 | ) 256 | print('filtered_field: ', type(filtered_field)) 257 | 258 | # projection 259 | # Limit output values to 0 ~ 1 260 | projected_field = mpa.tanh_projection(filtered_field, beta, eta) 261 | 262 | # interpolate to actual materials 263 | return projected_field.flatten() 264 | 265 | 266 | # Create Block with same size as design region 267 | geometry = [ 268 | mp.Block( 269 | center=design_region.center, size=design_region.size, material=design_variables 270 | ), 271 | mp.Block( 272 | center=mp.Vector3(0, -Sy/2 + PDsize/2, 0), size=mp.Vector3(Sx, PDsize, 0), material=SiO2 273 | ), 274 | # Use when DTI is present 275 | mp.Block( 276 | center=mp.Vector3(-design_region_width/3, -Sy/2 + PDsize/2, 0), size=mp.Vector3(subpixelsize, PDsize, 0), material=SiPD 277 | ), 278 | mp.Block( 279 | center=mp.Vector3(0, -Sy/2 + PDsize/2, 0), size=mp.Vector3(subpixelsize, PDsize, 0), material=SiPD 280 | ), 281 | mp.Block( 282 | center=mp.Vector3(design_region_width/3, -Sy/2 + PDsize/2, 0), size=mp.Vector3(subpixelsize, PDsize, 0), material=SiPD 283 | ) 284 | ] 285 | 286 | 287 | 288 | # Meep simulation setup 289 | sim = mp.Simulation( 290 | cell_size=cell_size, 291 | boundary_layers=pml_layers, 292 | geometry=geometry, 293 | sources=source, 294 | default_material=Air, # Empty space 295 | resolution=resolution, 296 | k_point = mp.Vector3(0,0,0) # bloch boundary 297 | ) 298 | 299 | ############################################################################################################################### 300 | # ## 2. Optimization Environment 301 | 302 | 303 | #Monitor position and size settings (focal point) 304 | monitor_position_0, monitor_size_0 = mp.Vector3(-design_region_width/3, -Sy/2 + PDsize - 0.5/resolution), mp.Vector3(subpixelsize,0) 305 | monitor_position_1, monitor_size_1 = mp.Vector3(0, -Sy/2 + PDsize - 0.5/resolution), mp.Vector3(subpixelsize,0) 306 | monitor_position_2, monitor_size_2 = mp.Vector3(design_region_width/3, -Sy/2 + PDsize - 0.5/resolution), mp.Vector3(subpixelsize,0) 307 | 308 | 309 | 310 | # Get Fourier transform for the area of monitor_size at monitor_position through FourierFields 311 | FourierFields_0 = mpa.FourierFields(sim,mp.Volume(center=monitor_position_0,size=monitor_size_0),mp.Ez,yee_grid=True) 312 | 313 | FourierFields_1 = mpa.FourierFields(sim,mp.Volume(center=monitor_position_1,size=monitor_size_1),mp.Ez,yee_grid=True) 314 | 315 | FourierFields_2= mpa.FourierFields(sim,mp.Volume(center=monitor_position_2,size=monitor_size_2),mp.Ez,yee_grid=True) 316 | 317 | 318 | ob_list = [FourierFields_0, FourierFields_1, FourierFields_2,] 319 | 320 | 321 | 322 | 323 | fred = [] 324 | fgreen = [] 325 | fblue = [] 326 | # J : Objective function 327 | # Take the squared absolute value of the Ez component at the center of the monitor measured by FourierFields 328 | # [frequency index, moniter index] 329 | 330 | 331 | flag = 0 332 | def J_0(fields_0, fields_1, fields_2): 333 | red = npa.sum(npa.abs(fields_0[21:30,:]) **2) 334 | green = npa.sum(npa.abs(fields_1[11:20,:]) ** 2) 335 | blue = npa.sum(npa.abs(fields_2[1:10,:]) ** 2) 336 | 337 | 338 | 339 | redfactor = 1 340 | greenfactor = 1 341 | bluefactor = 1 342 | 343 | if isinstance(red, np.floating): 344 | red_ = red 345 | else: 346 | red_ = red._value 347 | 348 | if isinstance(green, np.floating): 349 | green_ = green 350 | else: 351 | green_ = green._value 352 | 353 | if isinstance(blue, np.floating): 354 | blue_ = blue 355 | else: 356 | blue_ = blue._value 357 | 358 | 359 | 360 | fred.append(red_/redfactor) 361 | fgreen.append(green_/greenfactor) 362 | fblue.append(blue_/bluefactor) 363 | 364 | return blue/bluefactor + green/greenfactor + red/redfactor 365 | 366 | 367 | # Optimization settings 368 | opt = mpa.OptimizationProblem( 369 | simulation=sim, 370 | objective_functions=[J_0], 371 | objective_arguments=ob_list, 372 | design_regions=[design_region], 373 | frequencies=frequencies, 374 | decay_by=1e-3, # Ratio of field values remaining in monitor 375 | ) 376 | 377 | 378 | evaluation_history = [] 379 | cur_iter = [0] 380 | numevl = 0 381 | 382 | def f(v, gradient, beta): 383 | global numevl 384 | print("Current iteration: {}".format(cur_iter[0] + 1)) 385 | print("x: ", type(v)) 386 | x_image = x.reshape(Nx, Ny,1) 387 | print(x_image.shape) 388 | 389 | f0, dJ_du = opt([mapping(v, eta_i, beta)]) # compute objective and gradient 390 | print("FoM: ", f0) 391 | # f0, dJ_du = opt() 392 | adjoint_gradient = np.sum(dJ_du, axis =1 ) 393 | reshaped_gradients = adjoint_gradient.reshape(adjoint_gradient.shape[0], -1) 394 | 395 | # Computing the norm of each reshaped row 396 | norms = np.linalg.norm(reshaped_gradients, axis=1) 397 | 398 | # Calculating the mean of the norms 399 | adjgrad_norm = norms.mean() 400 | 401 | wandb.log({ 402 | "generated": [wandb.Image(1-x_image, caption=f"Iteration {numevl}")], 403 | 'fom ': f0, 404 | 'adjgrad norm': adjgrad_norm 405 | }, step = numevl), 406 | 407 | # Adjoint gradient 408 | if gradient.size > 0: 409 | gradient[:] = tensor_jacobian_product(mapping, 0)( 410 | v, eta_i, beta, np.sum(dJ_du, axis=1) 411 | ) # backprop 412 | 413 | evaluation_history.append(np.real(f0)) 414 | 415 | np.savetxt(design_dir+"structure_0"+str(numevl) +".txt", design_variables.weights) 416 | 417 | numevl += 1 418 | 419 | cur_iter[0] = cur_iter[0] + 1 420 | 421 | print("First FoM: {}".format(evaluation_history[0])) 422 | print("Current FoM: {}".format(np.real(f0))) 423 | 424 | 425 | return np.real(f0) 426 | 427 | ############################################################################################################################### 428 | 429 | # ## 3. Algorithm select 430 | 431 | 432 | 433 | n = Nx * Ny # number of parameters 434 | 435 | # Initial guess - random initial starting value 436 | #x = np.random.uniform(0.3, 0.7, n) 437 | 438 | x = np.ones(n)*0.5 439 | 440 | # lower and upper bounds (upper bound: 1, lower bound: 0) 441 | lb = np.zeros((Nx * Ny,)) 442 | ub = np.ones((Nx * Ny,)) 443 | 444 | # Optimization parameter 445 | cur_beta = 2 446 | beta_scale = 2 447 | num_betas = 5 448 | update_factor = 20 # number of iterations between beta updates 449 | ftol = 1e-5 450 | 451 | 452 | wandb.run.name = "min"+str(minimum_length)+"_"+"eta_i"+str(eta_i)+"_"+str(algorithm_name(algorithm))+"_"+"num_betas"+str(num_betas) 453 | 454 | 455 | for iters in range(num_betas): 456 | print("current beta: ", cur_beta) 457 | solver = nlopt.opt(algorithm, n) 458 | solver.set_lower_bounds(lb) # lower bounds 459 | solver.set_upper_bounds(ub) # upper bounds 460 | if cur_beta >=2**(num_betas+1): # 2^5 461 | solver.set_max_objective(lambda a, g: f(a, g, mp.inf)) 462 | solver.set_maxeval(1) # Set the maximum number of function evaluations 463 | else: 464 | solver.set_max_objective(lambda a, g: f(a, g, cur_beta)) 465 | solver.set_maxeval(update_factor) # Set the maximum number of function evaluations 466 | #solver.set_ftol_rel(ftol) # Set the relative tolerance for convergence 467 | x[:] = solver.optimize(x) 468 | print(x) 469 | x_ = copy.copy(x) 470 | cur_beta = cur_beta * beta_scale # Update the beta value for the next iteration 471 | wandb.log({ 472 | 'cur_beta': cur_beta 473 | }, step = numevl) 474 | wandb.config.update( {'update_factor': update_factor, 475 | 'beta_scale': beta_scale, 476 | 'num_betas': num_betas}) 477 | 478 | wandb.config.update({'algorithm': algorithm_name(algorithm)}) 479 | 480 | ############################################################################################################################### 481 | 482 | 483 | # ## 4. Save result 484 | 485 | #np.save("evaluation", evaluation_history) 486 | np.savetxt(design_dir+"evaluation.txt", evaluation_history) 487 | 488 | # FoM plot 489 | 490 | plt.figure() 491 | 492 | plt.plot(evaluation_history, "k-") 493 | plt.grid(False) 494 | plt.tick_params(axis='x', direction='in', pad = 5) 495 | plt.tick_params(axis='y', direction='in', pad = 10) 496 | plt.xlabel("Iteration") 497 | plt.ylabel("FoM") 498 | plt.savefig(design_dir+"FoMresult.png") 499 | plt.cla() # clear the current axes 500 | plt.clf() # clear the current figure 501 | plt.close() # closes the current figure 502 | 503 | 504 | def extract_elements(lst): 505 | # Create a list to store the results. 506 | result = [] 507 | 508 | # Iterate through the length of the list, extracting elements at indices that are multiples of 5. 509 | for i in range(0, len(lst), 5): 510 | result.append(lst[i]) 511 | 512 | return result 513 | 514 | # RGB FoM plot 515 | 516 | fred = extract_elements(fred) 517 | fgreen = extract_elements(fgreen) 518 | fblue = extract_elements(fblue) 519 | 520 | 521 | columns = ['red', 'green', 'blue'] 522 | wandb.log({"Intensity": wandb.plot.line_series( 523 | xs = range(len(fred)), 524 | ys = [fred, fgreen, fblue], 525 | keys = columns, 526 | xname = "Step" 527 | )}) 528 | 529 | plt.figure() 530 | 531 | plt.plot(fred, "r-") 532 | plt.plot(fgreen, "g-") 533 | plt.plot(fblue, "b-") 534 | plt.grid(False) 535 | plt.tick_params(axis='x', direction='in', pad = 5) 536 | plt.tick_params(axis='y', direction='in', pad = 10) 537 | plt.xlabel("Iteration") 538 | plt.ylabel("FoM") 539 | plt.savefig(design_dir+"FoMresultr.png") 540 | plt.cla() # clear the current axes 541 | plt.clf() # clear the current figure 542 | plt.close() # closes the current figure 543 | 544 | 545 | 546 | 547 | 548 | # Last design plot 549 | 550 | plt.imshow(design_variables.weights.reshape(Nx, Ny), cmap='binary') 551 | plt.colorbar() 552 | 553 | x_final = mapping(x, eta_i, mp.inf) 554 | 555 | f0, dJ_du = opt([x_final]) 556 | 557 | 558 | 559 | np.savetxt(design_dir+"lastdesign.txt", design_variables.weights) 560 | np.save(design_dir+"lastdesign.npy", design_variables.weights) 561 | wandb.log({ 562 | "generated_final": [wandb.Image(1-design_variables.weights.reshape(Nx, Ny))], 563 | 'final_fom ': np.real(f0), 564 | }) 565 | 566 | plt.savefig(design_dir+"lastdesign.png") 567 | plt.cla() # clear the current axes 568 | plt.clf() # clear the current figure 569 | plt.close() # closes the current figure 570 | 571 | 572 | sensitivity= np.sum(dJ_du, axis =1 ) 573 | 574 | fig = plt.figure(figsize=(20,20)) 575 | plt.imshow(np.squeeze(np.abs(sensitivity.reshape(Nx, Ny))))#(np.rot90(np.squeeze(np.abs(sensitivity[0].reshape(Nx, Ny))))); 576 | plt.xlabel("x") 577 | plt.ylabel("y") 578 | plt.savefig("sensitivity.png") 579 | cmap = plt.cm.get_cmap() 580 | colormapping = plt.cm.ScalarMappable(cmap=cmap) 581 | cbar = fig.colorbar(colormapping, ax=plt.gca()) 582 | wandb.log({"sensitivity": plt}) 583 | 584 | 585 | 586 | 587 | x_final = x_final.reshape(Nx, Ny) 588 | 589 | island_deleted = delete_islands_with_size_1(x_final.copy()) 590 | min_feature_size, min_feature_label, labeled_array = find_minimum_feature_size(island_deleted) 591 | 592 | fom_island, _ = opt([ 593 | island_deleted.flatten()]) 594 | print('min feature size: ', min_feature_size) 595 | print('min feature label: ', min_feature_label) 596 | print('labeled array: ', labeled_array) 597 | print('island_deleted: ', island_deleted.shape) 598 | 599 | #plt1 = highlight_minimum_island(island_deleted, min_feature_label, labeled_array, title='Minimum Feature 1s') 600 | 601 | wandb.log({ 602 | 'fom (final, island deleted1)': fom_island, 603 | "generated_final, island deleted1": [wandb.Image(1-island_deleted)], 604 | # "highlighted, island deleted1": [wandb.Image(plt1, caption=f"generated_final_island_deleted_highlghted")] 605 | }) 606 | array_converted_processed = delete_islands_with_size_1(1-island_deleted.copy()) 607 | array_original = 1- array_converted_processed 608 | 609 | min_feature_size2, min_feature_label2, labeled_array2 = find_minimum_feature_size(array_converted_processed) 610 | plt2 = highlight_minimum_island(array_original, min_feature_label2, labeled_array2) 611 | euler_number_original = euler_number(array_original, connectivity=1) 612 | _, num_islands1 = label_(island_deleted, connectivity=1, return_num=True) 613 | num_holes1 = num_islands1 - euler_number_original 614 | 615 | euler_number_converted = euler_number(array_converted_processed, connectivity=1) 616 | _, num_islands2 = label_(array_converted_processed, connectivity=1, return_num=True) 617 | num_holes2 = num_islands2 - euler_number_converted 618 | 619 | 620 | fom_island2, _ = opt([ 621 | array_original.flatten()] 622 | ) 623 | wandb.log({ 624 | "highlighted, island deleted2": [wandb.Image(plt2)], 625 | 626 | }) 627 | plt3 = highlight_minimum_island(array_original, min_feature_label, labeled_array) 628 | 629 | 630 | wandb.log({ 631 | 'fom (final, island deleted2)': fom_island2, 632 | "generated_final, island deleted2": [wandb.Image(1-array_original)], 633 | "highlighted, final": [wandb.Image(plt3)] 634 | }) 635 | 636 | print("array 0,0: ", array_original[0,0]) 637 | 638 | 639 | wandb.log({ 640 | 'mfs (original)': min_feature_size, 641 | 'mfs (converted)': min_feature_size2, 642 | 'number of islands (original)': num_islands1, 643 | 'number of islands (converted)':num_islands2, 644 | 'number of holes (original)' : num_holes1, 645 | 'number of holes (converted)': num_holes2, 646 | 'euler_number with 1 connectivity (original)': euler_number_original, 647 | 'euler_number with 1 connectivity (converted)': euler_number_converted, 648 | }) 649 | 650 | fig2 = plt.figure(figsize=(20, 8)) 651 | 652 | intensities_R = np.abs(opt.get_objective_arguments()[0][:,0]) ** 2 653 | plt.subplot(1,3,1) 654 | plt.plot(wavelengths/um_scale, intensities_R, "-o") 655 | plt.grid(True) 656 | # plt.xlim(0.38, 0.78) 657 | plt.xlabel("Wavelength (μm)") 658 | plt.ylabel("|Ez|^2 intensity (a.u.)") 659 | # plt.show() 660 | 661 | 662 | intensities_G = np.abs(opt.get_objective_arguments()[1][:,1]) ** 2 663 | fig3 = plt.figure(figsize=(20, 8)) 664 | plt.subplot(1,3,2) 665 | plt.plot(wavelengths/um_scale, intensities_G, "-o") 666 | plt.grid(True) 667 | # plt.xlim(0.38, 0.78) 668 | plt.xlabel("Wavelength (μm)") 669 | plt.ylabel("|Ez|^2 intensity (a.u.)") 670 | # plt.show() 671 | 672 | 673 | intensities_B = np.abs(opt.get_objective_arguments()[2][:,1]) ** 2 674 | plt.subplot(1,3,3) 675 | plt.plot(wavelengths/um_scale, intensities_B, "-o") 676 | plt.grid(True) 677 | # plt.xlim(0.38, 0.78) 678 | plt.xlabel("Wavelength (μm)") 679 | plt.ylabel("|Ez|^2 intensity (a.u.)") 680 | # plt.show() 681 | plt.savefig(design_dir+"FinalEz.png") 682 | wandb.log({"FinalEz": plt}) 683 | 684 | plt.cla() # clear the current axes 685 | plt.clf() # clear the current figure 686 | plt.close() # closes the current figure 687 | --------------------------------------------------------------------------------