diff --git a/farfield_testing.py b/farfield_testing.py index 8fd18f5d5..f4c49eb63 100755 --- a/farfield_testing.py +++ b/farfield_testing.py @@ -8,6 +8,7 @@ from autograd import grad, value_and_grad from autograd.test_util import check_grads from autograd.tracer import getval +from scipy.ndimage import gaussian_filter import tidy3d as td from tidy3d.web import run @@ -31,20 +32,20 @@ def get_spatial_coords_dict(simulation: td.Simulation, monitor: td.Monitor, fiel return coords -def make_field_data(x, sim, monitor, rng) -> td.FieldData: +def make_field_data_random(x, sim, monitor, rng) -> td.FieldData: field_cmps = {} grid = sim.discretize_monitor(monitor) for field_name in monitor.fields: coords = get_spatial_coords_dict(sim, monitor, field_name) coords["f"] = list(monitor.freqs) - data_shape = [len(coords[k]) for k in td.ScalarFieldDataArray._dims] + data_shape = np.array([len(coords[k]) for k in td.ScalarFieldDataArray._dims]) + dims = np.where(data_shape > 1)[0] - x = x[:, None] - s = (1, np.prod(data_shape)) - re = np.mean(x @ rng.uniform(-1, 1, s), axis=0) - im = np.mean(x @ rng.uniform(-1, 1, s), axis=0) - data = np.reshape(re + 1j * im, data_shape) + re = gaussian_filter(rng.uniform(-10, 10, min(data_shape[dims])), 1) + im = gaussian_filter(rng.uniform(-10, 10, min(data_shape[dims])), 1) + data = np.outer(x[: x.size // 2] ** 2, re) + 1j * np.outer(x[x.size // 2 :] ** 2, im) + data = np.reshape(data, data_shape) field_cmps[field_name] = td.ScalarFieldDataArray(data, coords=coords) @@ -57,11 +58,42 @@ def make_field_data(x, sim, monitor, rng) -> td.FieldData: ) -def make_sim_data(x, sim, seed): - rng = np.random.default_rng(seed) - data = [ - make_field_data(x, sim, mnt, rng) for mnt in sim.monitors if type(mnt) is td.FieldMonitor - ] +def make_field_data_from_file(x, monitor, sim_data) -> td.FieldData: + field_cmps = {} + field_data = sim_data[monitor.name] + for c in field_data.field_components: + field = getattr(field_data, c) + f = field.values + data = x[0] * np.real(f) + 1j * x[1] * np.imag(f) + field_cmps[c] = td.ScalarFieldDataArray(data, coords=field.coords) + + return td.FieldData( + monitor=field_data.monitor, + symmetry=field_data.symmetry, + symmetry_center=field_data.symmetry_center, + grid_expanded=field_data.grid_expanded, + **field_cmps, + ) + + +def make_sim_data(x, sim=None, seed=None, fp=None): + if fp is None and sim is not None: + rng = np.random.default_rng(seed) + data = [ + make_field_data_random(x, sim, mnt, rng, fp) + for mnt in sim.monitors + if type(mnt) is td.FieldMonitor + ] + elif fp is not None: + sim_data = td.SimulationData.from_file(fp) + sim = sim_data.simulation + data = [ + make_field_data_from_file(x, mnt, sim_data) + for mnt in sim.monitors + if type(mnt) is td.FieldMonitor + ] + else: + raise ValueError("yeah no") return td.SimulationData(simulation=sim, data=data) @@ -143,7 +175,7 @@ def objective(x, *, is_2d=False, local=False, emulated=False, seed=None): fp = f'autograd_projection_dbg_{"2d" if is_2d else "3d"}.hdf5' if emulated: # static sim, traced sim data sim = make_simulation(1.0, is_2d=is_2d) - sim_data = make_sim_data(x, sim, seed or 3467643734) + sim_data = make_sim_data(x, sim, fp="simulation_data.hdf5") elif Path(fp).is_file() and local: sim = make_simulation(x, is_2d=is_2d) sim_data = td.SimulationData.from_file(fp) @@ -169,9 +201,9 @@ def objective(x, *, is_2d=False, local=False, emulated=False, seed=None): def main(): seed = 46435543 np.random.seed(seed) # global seed for autograd - x = np.random.uniform(-1, 1, 10) + x = np.random.uniform(-10, 10, 2 * 71) - # objective(x, emulated=True, seed=seed) + # print(grad(objective)(x, emulated=True, seed=seed)) check_grads(objective, modes=["rev"], order=1)(x, emulated=True, seed=seed) diff --git a/gradcheck_nb.py b/gradcheck_nb.py new file mode 100755 index 000000000..d16a0744b --- /dev/null +++ b/gradcheck_nb.py @@ -0,0 +1,312 @@ +#!/usr/bin/env -S poetry run python +# ruff: noqa: F401 + +import autograd as ag +import autograd.numpy as anp +import matplotlib.pyplot as plt +import numpy as np +import optax +from autograd.test_util import check_grads + +import tidy3d as td +from farfield_testing import make_sim_data +from tidy3d import web +from tidy3d.web import run as run_adj + +td.config.logging_level = "ERROR" + +# 1 nanometer in units of microns (for conversion) +nm = 1e-3 + +# free space central wavelength +wavelength = 850 * nm + +# desired numerical aperture +# NA = 0.94 +NA = 0.03 + +# shape parameters of metalens unit cell (um) (refer to image above and see paper for details) +H = 430 * nm +S = 320 * nm + +# space between bottom PML and substrate (-z) +space_below_sub = 1 * wavelength + +# distance from unit cell top (z) to near field monitor +space_to_near_field = 1 * wavelength + +# thickness of substrate between source and Si unit cells +thickness_sub = 100 * nm + +# side length of entire metalens (um) +side_length = 12 + +sim_buffer_xy = 2 * wavelength + +# Number of unit cells in each x and y direction (NxN grid) +N = int(side_length / S) + +print(f"for diameter of {side_length:.1f} um, have {N} cells per side") +print(f"full metalens has area of {side_length**2:.1f} um^2 and {N*N} total cells") + +# Define material properties at 600 nm +n_Si = 3.84 +n_SiO2 = 1.46 +air = td.Medium(permittivity=1.0) +SiO2 = td.Medium(permittivity=n_SiO2**2) +Si = td.Medium(permittivity=n_Si**2) + +# define symmetry +symmetry = (-1, 1, 0) + +# using the wavelength in microns, one can use td.C_0 (um/s) to get frequency in Hz +# wavelength_meters = wavelength * meters +f0 = td.C_0 / wavelength + +# Compute the domain size in x, y (note: round down from side_length) +length_xy = N * S + +# focal length given diameter and numerical aperture +focal_length = length_xy / 2 / NA * np.sqrt(1 - NA**2) +print(f"{focal_length=}") + +# total domain size in z: (space -> substrate -> unit cell -> 1.7 focal lengths) +length_z = space_below_sub + thickness_sub + H + +# construct simulation size array +sim_size = (length_xy + 2 * sim_buffer_xy, length_xy + 2 * sim_buffer_xy, length_z) + +# define substrate +substrate = td.Structure( + geometry=td.Box.from_bounds( + rmin=(-td.inf, -td.inf, -1000), + rmax=(+td.inf, +td.inf, 0), + ), + medium=SiO2, +) + +# define coordinates of each unit cell +centers_x = S * np.arange(N) - length_xy / 2.0 + S / 2.0 +centers_y = S * np.arange(N) - length_xy / 2.0 + S / 2.0 +center_z = substrate.geometry.bounds[1][2] + H / 2 + +focal_z = center_z + H / 2 + focal_length + +x_centers, y_centers = np.meshgrid(centers_x, centers_y, indexing="ij") +xs = x_centers.flatten() +ys = y_centers.flatten() + + +def get_sizes(params): + """Returns the actual side lengths of the boxes as a function of design parameters from (-inf, +inf).""" + return S * (anp.tanh(params) + 1.0) / 2.0 + + +# initially, start with parameters of 0 (all boxes have side length S/2) +params0 = np.zeros(x_centers.shape) + + +def make_structures(params, apply_symmetry: bool = True): + """Make the Structure objects that will be used as .input_structures.""" + + sizes = get_sizes(params) + nx, ny = sizes.shape + geometries = [] + + for i in range(nx): + i_quad = max(i, nx - 1 - i) + for j in range(ny): + j_quad = max(j, ny - 1 - j) + size = sizes[i_quad, j_quad] + x0 = x_centers[i, j] + y0 = y_centers[i, j] + + if apply_symmetry and symmetry[0] != 0 and x0 < -S / 2: + continue + + if apply_symmetry and symmetry[1] != 0 and y0 < -S / 2: + continue + + # geometry = td.Box(center=(x0, y0, center_z), size=(size, size, H)) + geometry = td.Cylinder(center=(x0, y0, center_z), length=H, radius=size / 2) + + geometries.append(geometry) + geo_group = td.GeometryGroup(geometries=geometries) + medium = td.Medium(permittivity=n_Si**2) + + return [td.Structure(medium=medium, geometry=geo_group)] + # return [td.Structure(medium=medium, geometry=geo) for geo in geometries] + + +structures = make_structures(params0) + +# steps per unit cell along x and y +grids_per_unit_length = 10 + +# uniform mesh in x and y +grid_x = td.UniformGrid(dl=S / grids_per_unit_length) +grid_y = td.UniformGrid(dl=S / grids_per_unit_length) + +# in z, use an automatic nonuniform mesh with the wavelength being the "unit length" +grid_z = td.AutoGrid(min_steps_per_wvl=grids_per_unit_length) + +# we need to supply the wavelength because of the automatic mesh in z +grid_spec = td.GridSpec(wavelength=wavelength, grid_x=grid_x, grid_y=grid_y, grid_z=grid_z) + +# put an override box over the pillars to avoid parsing a large amount of structures in the mesher +grid_spec = grid_spec.copy( + update=dict( + override_structures=[ + td.Structure( + geometry=td.Box.from_bounds( + rmin=(-td.inf, -td.inf, -length_z / 2 + space_below_sub), + rmax=(td.inf, td.inf, center_z + H / 2), + ), + medium=Si, + ) + ] + ) +) + +# Bandwidth in Hz +fwidth = f0 / 10.0 + +# time dependence of source +gaussian = td.GaussianPulse(freq0=f0, fwidth=fwidth, phase=0) + +source = td.PlaneWave( + source_time=gaussian, + size=(td.inf, td.inf, 0), + center=(0, 0, -0.1), + direction="+", + pol_angle=0, +) + +run_time = 50 / fwidth + +monitor_near = td.FieldMonitor( + center=[0.0, 0.0, H + 0.1], + size=[td.inf, td.inf, 0], + freqs=[f0], + name="near_field", + colocate=False, +) + + +def make_sim(angles, apply_symmetry: bool = True): + metalens = make_structures(angles, apply_symmetry=apply_symmetry) + sim = td.Simulation( + size=sim_size, + center=(0, 0, H / 2), + grid_spec=grid_spec, + structures=[substrate] + metalens, + sources=[source], + monitors=[monitor_near], + run_time=run_time, + boundary_spec=td.BoundarySpec( + x=td.Boundary.absorber(), y=td.Boundary.absorber(), z=td.Boundary.pml() + ), + symmetry=symmetry, + ) + return sim + + +def measure_focal_intensity(sim_data: td.SimulationData) -> float: + """Measures electric intensity at focal point.""" + + monitor_far = td.FieldProjectionAngleMonitor( + center=(0, 0, focal_z), + size=(0, 0, 0), + freqs=[f0], + phi=[0], + theta=[0], + proj_distance=100, + far_field_approx=True, + name="far_field", + ) + + # first, we construct the classes which compute far fields locally on your machine + n2f = td.FieldProjector.from_near_field_monitors( + sim_data=sim_data, + near_monitors=[monitor_near], # only supply the non-downsampled surface monitors as sources + normal_dirs=["+"], + pts_per_wavelength=10, + ) + + far_fields = n2f.project_fields(monitor_far) + # return anp.sum(anp.abs(far_fields.Etheta.values)) + return abs(far_fields.power.values.item()) + + +def J(params) -> float: + """Objective function, returns intensity at focal point as a function of params.""" + sim = make_sim(params) + # sim = make_sim(np.zeros(x_centers.shape)) + sim_data = run_adj(sim, task_name="metalens_invdes", verbose=True) + # sim_data = make_sim_data(params, fp="simulation_data.hdf5") + # sim_data = td.SimulationData.from_file("simulation_data.hdf5") + return measure_focal_intensity(sim_data) + + +# val = J(params0) + +# # params0 = np.random.uniform(-10, 10, 2 * 452) +# # params0 = np.array([1.0, 1.0]) +# # print(J(params0)) +dJ = ag.value_and_grad(J) +val, grad = dJ(params0) +print(val, grad) +# exit() + +# check_grads(J, modes=["rev"], order=1)(params0) +# exit() + +params_empty = -1e5 * np.ones_like(params0) +J_empty = np.array(J(params_empty)) + + +def J_normalized(params): + return J(params) / J_empty + + +val_normalized = val / J_empty + +dJ_normalized = ag.value_and_grad(J_normalized) + +# hyperparameters +num_steps = 10 +learning_rate = 1e-2 + +# initialize adam optimizer with starting parameters +params = np.array(params0) +optimizer = optax.adam(learning_rate=learning_rate) +opt_state = optimizer.init(params) + +# store history +J_history = [val_normalized] +params_history = [params0] + +for i in range(num_steps): + # compute gradient and current objective function value + value, gradient = dJ_normalized(params) + + # outputs + print(f"step = {i + 1}") + print(f"\tJ = {value:.4e}") + print(f"\tgrad_norm = {np.linalg.norm(gradient):.4e}") + + # compute and apply updates to the optimizer based on gradient (-1 sign to maximize obj_fn) + updates, opt_state = optimizer.update(gradient, opt_state, params) + params = optax.apply_updates(params, updates) + params = np.array(params) + + # save history + J_history.append(value) + params_history.append(params) + +params_after = params_history[-1] + +plt.plot(J_history) +plt.xlabel("iterations") +plt.ylabel("objective function (focusing intensity enhancement)") +plt.show() diff --git a/nb.py b/nb.py new file mode 100755 index 000000000..9ff7c2ad6 --- /dev/null +++ b/nb.py @@ -0,0 +1,310 @@ +#!/usr/bin/env -S poetry run python +# ruff: noqa: F401 + +import autograd as ag +import autograd.numpy as anp +import matplotlib.pyplot as plt +import numpy as np +import optax +from autograd.test_util import check_grads + +import tidy3d as td +from tidy3d import web +from tidy3d.web import run as run_adj + +td.config.logging_level = "ERROR" + +# 1 nanometer in units of microns (for conversion) +nm = 1e-3 + +# free space central wavelength +wavelength = 850 * nm + +# desired numerical aperture +# NA = 0.94 +NA = 0.05 + +# shape parameters of metalens unit cell (um) (refer to image above and see paper for details) +H = 430 * nm +S = 320 * nm + +# space between bottom PML and substrate (-z) +space_below_sub = 1 * wavelength + +# distance from unit cell top (z) to near field monitor +space_to_near_field = 1 * wavelength + +# thickness of substrate between source and Si unit cells +thickness_sub = 100 * nm + +# side length of entire metalens (um) +side_length = 12 + +sim_buffer_xy = 2 * wavelength + +# Number of unit cells in each x and y direction (NxN grid) +N = int(side_length / S) + +print(f"for diameter of {side_length:.1f} um, have {N} cells per side") +print(f"full metalens has area of {side_length**2:.1f} um^2 and {N*N} total cells") + +# Define material properties at 600 nm +n_Si = 3.84 +n_SiO2 = 1.46 +air = td.Medium(permittivity=1.0) +SiO2 = td.Medium(permittivity=n_SiO2**2) +Si = td.Medium(permittivity=n_Si**2) + +# define symmetry +symmetry = (-1, 1, 0) + +# using the wavelength in microns, one can use td.C_0 (um/s) to get frequency in Hz +# wavelength_meters = wavelength * meters +f0 = td.C_0 / wavelength + +# Compute the domain size in x, y (note: round down from side_length) +length_xy = N * S + +# focal length given diameter and numerical aperture +focal_length = length_xy / 2 / NA * np.sqrt(1 - NA**2) +print(f"{focal_length=}") + +# total domain size in z: (space -> substrate -> unit cell -> 1.7 focal lengths) +length_z = space_below_sub + thickness_sub + H + +# construct simulation size array +sim_size = (length_xy + 2 * sim_buffer_xy, length_xy + 2 * sim_buffer_xy, length_z) + +# define substrate +substrate = td.Structure( + geometry=td.Box.from_bounds( + rmin=(-td.inf, -td.inf, -1000), + rmax=(+td.inf, +td.inf, 0), + ), + medium=SiO2, +) + +# define coordinates of each unit cell +centers_x = S * np.arange(N) - length_xy / 2.0 + S / 2.0 +centers_y = S * np.arange(N) - length_xy / 2.0 + S / 2.0 +center_z = substrate.geometry.bounds[1][2] + H / 2 + +focal_z = center_z + H / 2 + focal_length + +x_centers, y_centers = np.meshgrid(centers_x, centers_y, indexing="ij") +xs = x_centers.flatten() +ys = y_centers.flatten() + + +def get_sizes(params): + """Returns the actual side lengths of the boxes as a function of design parameters from (-inf, +inf).""" + return S * (anp.tanh(params) + 1.0) / 2.0 + + +# initially, start with parameters of 0 (all boxes have side length S/2) +params0 = np.ones(x_centers.shape) + + +def make_structures(params, apply_symmetry: bool = True): + """Make the Structure objects that will be used as .input_structures.""" + + coords = dict(x=centers_x, y=centers_y, z=[0]) + eps = 2 + anp.tanh(anp.reshape(params, (*params.shape, 1))) + permittivity = td.SpatialDataArray(eps, coords=coords) + custom_medium = td.CustomMedium(permittivity=permittivity) + box = td.Box(center=(0, 0, center_z), size=(length_xy, length_xy, H)) + return [td.Structure(geometry=box, medium=custom_medium)] + + sizes = get_sizes(params) + nx, ny = sizes.shape + geometries = [] + + for i in range(nx): + i_quad = max(i, nx - 1 - i) + for j in range(ny): + j_quad = max(j, ny - 1 - j) + size = sizes[i_quad, j_quad] + x0 = x_centers[i, j] + y0 = y_centers[i, j] + + if apply_symmetry and symmetry[0] != 0 and x0 < -S / 2: + continue + + if apply_symmetry and symmetry[1] != 0 and y0 < -S / 2: + continue + + geometry = td.Cylinder(center=(x0, y0, center_z), length=H, radius=size / 2) + + geometries.append(geometry) + geo_group = td.GeometryGroup(geometries=geometries) + medium = td.Medium(permittivity=n_Si**2) + + return [td.Structure(medium=medium, geometry=geo_group)] + + +structures = make_structures(params0) + +# steps per unit cell along x and y +grids_per_unit_length = 10 + +# uniform mesh in x and y +grid_x = td.UniformGrid(dl=S / grids_per_unit_length) +grid_y = td.UniformGrid(dl=S / grids_per_unit_length) + +# in z, use an automatic nonuniform mesh with the wavelength being the "unit length" +grid_z = td.AutoGrid(min_steps_per_wvl=grids_per_unit_length) + +# we need to supply the wavelength because of the automatic mesh in z +grid_spec = td.GridSpec(wavelength=wavelength, grid_x=grid_x, grid_y=grid_y, grid_z=grid_z) + +# put an override box over the pillars to avoid parsing a large amount of structures in the mesher +grid_spec = grid_spec.copy( + update=dict( + override_structures=[ + td.Structure( + geometry=td.Box.from_bounds( + rmin=(-td.inf, -td.inf, -length_z / 2 + space_below_sub), + rmax=(td.inf, td.inf, center_z + H / 2), + ), + medium=Si, + ) + ] + ) +) + +# Bandwidth in Hz +fwidth = f0 / 10.0 + +# time dependence of source +gaussian = td.GaussianPulse(freq0=f0, fwidth=fwidth, phase=0) + +source = td.PlaneWave( + source_time=gaussian, + size=(td.inf, td.inf, 0), + center=(0, 0, -0.1), + direction="+", + pol_angle=0, +) + +run_time = 50 / fwidth + +monitor_near = td.FieldMonitor( + center=[0.0, 0.0, H + 0.1], + size=[td.inf, td.inf, 0], + freqs=[f0], + name="near_field", + colocate=False, +) + + +def make_sim(angles, apply_symmetry: bool = True): + metalens = make_structures(angles, apply_symmetry=apply_symmetry) + sim = td.Simulation( + size=sim_size, + center=(0, 0, H / 2), + grid_spec=grid_spec, + structures=[substrate] + metalens, + sources=[source], + monitors=[monitor_near], + run_time=run_time, + boundary_spec=td.BoundarySpec( + x=td.Boundary.absorber(), y=td.Boundary.absorber(), z=td.Boundary.pml() + ), + symmetry=symmetry, + ) + return sim + + +def measure_focal_intensity(sim_data: td.SimulationData) -> float: + """Measures electric intensity at focal point.""" + + monitor_far = td.FieldProjectionAngleMonitor( + center=(0, 0, focal_z), + size=(0, 0, 0), + freqs=[f0], + phi=[0], + theta=[0], + proj_distance=focal_length, + far_field_approx=True, + name="far_field", + ) + + # first, we construct the classes which compute far fields locally on your machine + n2f = td.FieldProjector.from_near_field_monitors( + sim_data=sim_data, + near_monitors=[monitor_near], # only supply the non-downsampled surface monitors as sources + normal_dirs=["+"], + pts_per_wavelength=10, + ) + + far_fields = n2f.project_fields(monitor_far) + return abs(far_fields.power.values.item()) + + +def J(params) -> float: + """Objective function, returns intensity at focal point as a function of params.""" + sim = make_sim(params) + sim_data = run_adj(sim, task_name="metalens_invdes_dbg", verbose=True) + return measure_focal_intensity(sim_data) + + +check_grads(J, modes=["rev"], order=1)(params0) +exit() + +val = J(params0) + +# dJ = ag.value_and_grad(J) +# val, grad = dJ(params0) +# print(val) +# print(grad) +# exit() + +params_empty = -1e5 * np.ones_like(params0) +J_empty = np.array(J(params_empty)) + + +def J_normalized(params): + return J(params) / J_empty + + +val_normalized = val / J_empty + +dJ_normalized = ag.value_and_grad(J_normalized) + +# hyperparameters +num_steps = 10 +learning_rate = 1e-2 + +# initialize adam optimizer with starting parameters +params = np.array(params0) +optimizer = optax.adam(learning_rate=learning_rate) +opt_state = optimizer.init(params) + +# store history +J_history = [val_normalized] +params_history = [params0] + +for i in range(num_steps): + # compute gradient and current objective function value + value, gradient = dJ_normalized(params) + + # outputs + print(f"step = {i + 1}") + print(f"\tJ = {value:.4e}") + print(f"\tgrad_norm = {np.linalg.norm(gradient):.4e}") + + # compute and apply updates to the optimizer based on gradient (-1 sign to maximize obj_fn) + updates, opt_state = optimizer.update(-gradient, opt_state, params) + params = optax.apply_updates(params, updates) + params = np.array(params) + + # save history + J_history.append(value) + params_history.append(params) + +params_after = params_history[-1] + +plt.plot(J_history) +plt.xlabel("iterations") +plt.ylabel("objective function (focusing intensity enhancement)") +plt.show() diff --git a/objarray_perf.py b/objarray_perf.py new file mode 100755 index 000000000..dc1d80733 --- /dev/null +++ b/objarray_perf.py @@ -0,0 +1,46 @@ +#!/usr/bin/env -S poetry run python +# ruff: noqa: F401 + +from time import perf_counter + +import autograd.numpy as np +import numpy as onp +from autograd import grad + + +def main(): + def f(x, cast=False): + for _ in range(10): + if cast: + x = onp.array(x) + x = np.array(x.tolist()) + for ii in range(10): + if ii % 2 == 0: + x = x - np.mean(x) + else: + x = x + np.mean(x) + return np.mean(x) + + df = grad(f) + + x = np.random.uniform(-1, 1, 100000) + + t0 = perf_counter() + f(x, False) + print(f"fwd, no cast: {perf_counter() - t0:.5f}s") + + t0 = perf_counter() + f(x, True) + print(f"fwd, cast: {perf_counter() - t0:.5f}s") + + t0 = perf_counter() + df(x, False) + print(f"bwd, no cast: {perf_counter() - t0:.5f}s") + + t0 = perf_counter() + df(x, True) + print(f"bwd, cast: {perf_counter() - t0:.5f}s") + + +if __name__ == "__main__": + main() diff --git a/run.py b/run.py new file mode 100755 index 000000000..2870b0b8a --- /dev/null +++ b/run.py @@ -0,0 +1,129 @@ +#!/usr/bin/env -S poetry run python +# ruff: noqa: F401 + +from importlib import reload +from pathlib import Path +from time import perf_counter + +import autograd.numpy as anp +import matplotlib.pyplot as plt +import numpy as np +from autograd import value_and_grad +from autograd.test_util import check_grads + +import tidy3d as td +from tidy3d.web import run + +# import tidy3d.web.api.webapi as webapi +# from tests.utils import run_emulated + +# fp = Path("field_projection.hdf5") + +# webapi.run = run_emulated +# # webapi.run = lambda *args, **kwargs: td.SimulationData.from_hdf5(fp) +# from tidy3d.web.api.autograd import autograd # noqa +# from tidy3d.web.api.autograd.autograd import run # noqa + +# reload(autograd) + + +wavelength = 0.75 +f0 = td.C_0 / wavelength + + +def make_sim(width=1.0): + # plate = td.Structure(geometry=td.Box(size=(td.inf, 0.2, td.inf)), medium=td.PECMedium()) + aperture = td.Structure( + geometry=td.Box(size=(width, 0.3, width)), medium=td.Medium(permittivity=2) + ) + + source = td.PlaneWave( + center=(0, -0.3, 0), + size=(td.inf, 0, td.inf), + source_time=td.GaussianPulse(freq0=f0, fwidth=f0 / 10), + direction="+", + pol_angle=np.pi / 2, + ) + + monitor_near = td.FieldMonitor( + center=(0, 0.3, 0), + size=(td.inf, 0, td.inf), + # freqs=[f0, f0 + 1], + freqs=[f0], + name="near_field", + colocate=False, + ) + + return td.Simulation( + size=(2, 2, 2), + grid_spec=td.GridSpec.auto(wavelength, min_steps_per_wvl=20), + # structures=(plate, aperture), + structures=(aperture,), + sources=[source], + monitors=[monitor_near], + run_time=1e-13, + boundary_spec=td.BoundarySpec.all_sides(boundary=td.PML()), + ) + + +def get_far_field_intensity(data: td.SimulationData): + sim = data.simulation + r_proj = 50 * wavelength + + xy_proj = np.linspace(-0.2, 2, 5) + monitor_far = td.FieldProjectionCartesianMonitor( + center=sim.monitors[0].center, + size=sim.monitors[0].size, + freqs=sim.monitors[0].freqs, + name="far_field", + x=xy_proj, + y=xy_proj, + proj_axis=1, + proj_distance=r_proj, + far_field_approx=True, + ) + + # theta_proj = np.linspace(np.pi / 10, np.pi - np.pi / 10, 2) + # phi_proj = np.linspace(np.pi / 10, np.pi - np.pi / 10, 3) + # monitor_far_sph = td.FieldProjectionAngleMonitor( + # center=sim.monitors[0].center, + # size=sim.monitors[0].size, + # freqs=sim.monitors[0].freqs, + # name="far_field", + # phi=tuple(phi_proj), + # theta=tuple(theta_proj), + # proj_distance=r_proj, + # far_field_approx=True, + # ) + + projector = td.FieldProjector.from_near_field_monitors( + sim_data=data, + near_monitors=[sim.monitors[0]], + normal_dirs=["+"], + pts_per_wavelength=10, + ) + + projected_fields = projector.project_fields(monitor_far) + return anp.mean(anp.abs(anp.array(projected_fields.fields_cartesian.Ez.values.tolist()))) + + +def objfun(param: float) -> float: + sim = make_sim(param) + data = run(sim, task_name="adjoint_quickstart", verbose=True) + return get_far_field_intensity(data) + + +t0 = perf_counter() +v0 = objfun(1.0) +print("fwd time: ", perf_counter() - t0) + +t0 = perf_counter() +v1, g = value_and_grad(objfun)(1.0) +print("adj time: ", perf_counter() - t0) + +print("fwd result: ", v0) +print("adj result: ", v1) + +print("grads: ", g) + +check_grads(objfun, modes=["rev"], order=1)(1.0) diff --git a/tidy3d/components/data/monitor_data.py b/tidy3d/components/data/monitor_data.py index 5c6ab53a6..27a9b6a6a 100644 --- a/tidy3d/components/data/monitor_data.py +++ b/tidy3d/components/data/monitor_data.py @@ -1099,6 +1099,8 @@ def shift_value(coords) -> float: field_component = field_component.sel(f=freq0) forward_amps = field_component.values values = -1j * forward_amps + if "H" in name: + values *= -1 coords = dict(field_component.coords.copy()) for dim, key in enumerate("xyz"): coords[key] = np.array(coords[key]) - source_geo.center[dim] diff --git a/tidy3d/web/api/autograd/autograd.py b/tidy3d/web/api/autograd/autograd.py index 2c5093ce0..1b148d1e1 100644 --- a/tidy3d/web/api/autograd/autograd.py +++ b/tidy3d/web/api/autograd/autograd.py @@ -732,9 +732,10 @@ def setup_adj( td.log.info("Running custom vjp (adjoint) pipeline.") # immediately filter out any data_vjps with all 0's in the data - data_fields_vjp = { - key: get_static(value) for key, value in data_fields_vjp.items() if not np.all(value == 0.0) - } + # data_fields_vjp = { + # key: get_static(value) for key, value in data_fields_vjp.items() if not np.all(value == 0.0) + # } + data_fields_vjp = {key: get_static(value) for key, value in data_fields_vjp.items()} # insert the raw VJP data into the .data of the original SimulationData sim_data_vjp = sim_data_orig.insert_traced_fields(field_mapping=data_fields_vjp)