Anisotropic material gradient support #1801

smartalecH · 2021-10-26T13:44:23Z

This needs to work before #1780 (since subpixel smoothing induces anisotropic tensor elements).

Essentially, this PR properly recombines the forward and adjoint fields when calculating the gradients, regardless of what the material looks like (e.g. with or without off-diagonal elements). To do so, we store the D fields rather than the E fields, which means we look at chi1inv rather than chi1.

The process is straightforward:

the adjoint fields are restricted to the respective epsilon nodes
the forward fields are interpolated to the same nodes
the two are "recombined" using the gradient of the epsilon inverse tensor at that point (i.e. a weighted inner product).

This is the proper interpolation/restriction scheme, as this is the gradient of what's actually implemented in the code.

As a nice side effect, this makes the gradients significantly more accurate (a few orders of magnitude when comparing relative errors for some initial tests).

TODO:

Fix merge conflicts
Fix failing tests
Renable support for dispersive/~~conductive~~ media (I turned this off in order to simplify the code while debugging)
Add test with material containing off-diagonal components (e.g. rotated sapphire)
Cleanup

NOTE: conductive media will require a separate PR (see #1811)

smartalecH · 2021-10-27T18:49:52Z

@ianwilliamson, I'm not quite sure what's going on with the test_adjoint_jax.py test suite. This PR passes the test_adjoint_solver.py suite just fine.

Could you run the test on your end and look into it a bit?

…l_grad_mod

ianwilliamson · 2021-10-27T20:58:40Z

@ianwilliamson, I'm not quite sure what's going on with the test_adjoint_jax.py test suite. This PR passes the test_adjoint_solver.py suite just fine.

Could you run the test on your end and look into it a bit?

From a quick glance, it looks like this is updating the API of mp._get_gradient() to use the D fields instead of the E fields? At a minimum, we most likely want to update the field components that are specified here:

meep/python/adjoint/utils.py

Line 9 in e8c2023

_ADJOINT_FIELD_COMPONENTS = [mp.Ex, mp.Ey, mp.Ez]

Are there any other APIs that are going to be changed as part of this?

Separately, it may be nice to orient towards reusing common functions between OptimizationProblem and the JAX wrapper. This has already happened for parts of the JAX wrapper here: https://github.com/NanoComp/meep/blob/master/python/adjoint/utils.py. So maybe the OptimizationProblem can be switched over to use the functions from this module? FWIW, it's all unit tested.

smartalecH · 2021-10-27T21:07:39Z

Oof -- I didn't realize we diverged. We certainly need to refactor and use the same codebase for both APIs.

In addition to the change in the field components, the fields themselves must be stored on the Yee grid (yee_grid=False). Currently, utils.py sets this to True... which is a bit puzzling because it uses the get_gradient() routine which requires the fields to be on the Yee grid...

Either way, this will require some significant refactoring, as each component has a different array size that must be stored and carefully processed.

(we should add cylindrical support too)

ianwilliamson · 2021-10-27T21:25:21Z

In addition to the change in the field components, the fields themselves must be stored on the Yee grid (yee_grid=False). Currently, utils.py sets this to True... which is a bit puzzling because it uses the get_gradient() routine which requires the fields to be on the Yee grid...

I think the setting in utils.py:

meep/python/adjoint/utils.py

Line 71 in 5c02ff2

yee_grid=True,

is consistent with what's in OptimizationProblem:

meep/python/adjoint/optimization_problem.py

Line 218 in 5c02ff2

yee_grid=True,

Am I missing something?

smartalecH · 2021-10-27T21:32:33Z

Am I missing something?

Ah you're right! Sorry for the confusion (I've been at this way too long...)

ianwilliamson · 2021-10-27T21:37:46Z

Am I missing something?

Ah you're right! Sorry for the confusion (I've been at this way too long...)

No worries.

I don't think this would be too bad to factor into OptimizationProblem. Most of the functions in utils.py should correspond to very specific blocks and sequences of operations performed in OptimizationProblem.

python/adjoint/optimization_problem.py

codecov-commenter · 2021-10-28T19:58:54Z

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

Attention: Patch coverage is 95.31250% with 3 lines in your changes missing coverage. Please review.

Project coverage is 73.52%. Comparing base (c65a149) to head (37174ef).
Report is 356 commits behind head on master.

Files with missing lines	Patch %	Lines
python/adjoint/optimization_problem.py	83.33%	2 Missing ⚠️
python/adjoint/utils.py	97.87%	1 Missing ⚠️

❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1801      +/-   ##
==========================================
- Coverage   73.92%   73.52%   -0.41%     
==========================================
  Files          17       17              
  Lines        4924     4876      -48     
==========================================
- Hits         3640     3585      -55     
- Misses       1284     1291       +7

Files with missing lines	Coverage Δ
python/adjoint/__init__.py	`83.33% <100.00%> (+1.51%)`	⬆️
python/adjoint/objective.py	`91.81% <100.00%> (+0.04%)`	⬆️
python/adjoint/wrapper.py	`98.48% <ø> (+7.03%)`	⬆️
python/adjoint/utils.py	`95.23% <97.87%> (+1.79%)`	⬆️
python/adjoint/optimization_problem.py	`56.54% <83.33%> (-13.55%)`	⬇️

... and 3 files with indirect coverage changes

…l_grad_mod

smartalecH · 2021-11-01T14:59:27Z

Should now be ready for review/merging.

@ianwilliamson, I ended up refactoring the two APIs quite a bit, such that all of the major calls are within utils.py (per your suggestion). I also removed a few redundant routines that have since moved to simulation.py (specifically the DFT convergence test, which is now much more efficient thanks to #1740). There's still more that could be done in the future (see #1805) but that's better suited for another PR -- I only refactored the relevant pieces to this PR in order to simplify reviewing. Let me know if you have any other suggestions.

As described in #1811, the gradients are still incorrect for materials containing conductivities. Resolving that is best suited for another PR since the issue existed before this PR.

oskooi · 2021-11-02T16:28:32Z

python/adjoint/utils.py

+            '''We need to correct for the rare cases that get_dft_array
+            returns a singleton element for the forward and adjoint fields.
+            This only occurs when we are in 2D and only working with a particular
+            polarization (as the other fields are never stored). Our get_gradient


Perhaps we can be more explicit here and state that in 2D, the E_z polarization consists of a single scalar E_z field.

oskooi · 2021-11-02T21:06:03Z

src/meepgeom.cpp

+    meep::volume v(r);
+    LOOP_OVER_DIRECTIONS(dim, d){
+      v.set_direction_min(d, r.in_direction(d)-sd/2);
+      v.set_direction_max(d, r.in_direction(d)+sd/2);


The smoothing diameter sd is in units of the voxel dimension Δx (=1/resolution) rather than in absolute units. This means sd/2 should sd/(2*Δx). This is similar to how it is set up in grid_volume::dV:

meep/src/vec.cpp

Line 628 in fc5ea4a

const double hinva = 0.5 * inva * diameter;

I was just copying what's performed during initialization:

meep/src/anisotropic_averaging.cpp

Line 204 in 3eabcef

const double smoothing_diameter = 1.0; // FIXME: make user-changable?

Am I misunderstanding how this is being used?

Ah I think I see. The code also uses gv.dv() to do some scaling under the hood. Got it.

oskooi · 2021-11-02T22:25:32Z

src/meepgeom.cpp

+            meep::vec eps2 = gv[ieps2];
+            cyl_scale = (gv.dim == meep::Dcyl) ? eps2.r() : 1;
+            prod = 0.5*adj*get_material_gradient(eps2, adjoint_c, forward_c, fwd_avg, frequencies[f_i], geps, 1e-6);
+            material_grids_addgradient_point(v+ng*f_i, vec_to_vector3(eps2), scalegrad*prod.real()*cyl_scale, geps);


There is redundant code here that can be separated into a function and called twice with different arguments (fwd_p, fwd_pf, eps1) and (fwd_pa, fwd_paf, eps2).

I thought so too... but everything worth refactoring is already refactored in get_material_gradient and material_grids_addgradient_point.

Pulling it out further just made it more unreadable (and requires a lot more parameters to pass than just those 3).

prod could technically be placed within material_grids_addgradient_point. cyl_scale is repeated, but again, refactoring actually makes it harder to read.

ianwilliamson · 2021-11-02T22:58:27Z

python/adjoint/optimization_problem.py

-                    self.d_E[nb][ic][f, :, :, :] = atleast_3d(
-                        self.sim.get_dft_array(dgm, c, f))
+        # Store forward fields for each set of design variables in array
+        self.d_E = utils.gather_design_region_fields(self.sim,self.design_region_monitors,self.frequencies)


This is the D field now, correct? Perhaps this var should be renamed?

ianwilliamson · 2021-11-02T22:58:47Z

python/adjoint/optimization_problem.py

-
-        if self.sim.is_cylindrical or self.sim.dimensions == mp.CYLINDRICAL:
+            # Store adjoint fields for each design set of design variables
+            self.a_E.append(utils.gather_design_region_fields(self.sim,self.design_region_monitors,self.frequencies))


This is also the D field? Rename the var?

ianwilliamson · 2021-11-02T23:02:08Z

python/adjoint/utils.py

+            s = sim.get_array_slice_dimensions(c, vol=self.volume)[0]
+            if (onp.squeeze(fields_a[ci][0,...]).size == 1):
+                fields_a[ci] = onp.zeros(onp.insert(s,0,num_freqs))
+                fields_f[ci] = onp.zeros(onp.insert(s,0,num_freqs))


It would be more readable to do onp.zeros((num_freqs,) + s). Also, would help to use a more descriptive name than s... maybe spatial_shape or something similar?

maybe spatial_shape or something similar?

Good idea

onp.zeros((num_freqs,) + s)

unfortunately s is a numpy array, and can't be joined like python lists. This seems pretty standard when prepending to numpy arrays, from what I've seen (there are hundreds of possibilities of course...)

ianwilliamson · 2021-11-02T23:06:27Z

python/adjoint/utils.py

+        shapes = []
+        com = _compute_components(sim)
+        scalegrad = 1
+        for ci, c in enumerate(com):


I suggest for component_idx, component for better readability

ianwilliamson · 2021-11-02T23:08:40Z

python/adjoint/utils.py

+        num_freqs = onp.array(frequencies).size
+        shapes = []
+        com = _compute_components(sim)
+        scalegrad = 1


Can you add a comment as to why scalegrad is needed? If it's hardcoded to 1 here, can it just be removed?

ianwilliamson · 2021-11-02T23:10:50Z

python/adjoint/utils.py

+            proper shape as the design_region.            
+            '''
+            s = sim.get_array_slice_dimensions(c, vol=self.volume)[0]
+            if (onp.squeeze(fields_a[ci][0,...]).size == 1):


Why do you need the squeeze? Can't you just check the size?

ianwilliamson · 2021-11-02T23:12:56Z

python/adjoint/utils.py

+    def get_gradient(self, sim, fields_a, fields_f, frequencies):
+        num_freqs = onp.array(frequencies).size
+        shapes = []
+        com = _compute_components(sim)


Maybe just move the call to _compute_components() inside enumerate(...)

… diameter

* init * get everything compiled and default passing tests * new start * anisotropic gradients are working * fix multifrequency bug and other cleanup * add support for cyl coordindates * add rotated sapphire to adjoint solver test * get tests passing locally * enable dispersive material support * some minor cleanup * clean up comments and readability issues and fix scaling of smoothing diameter Co-authored-by: Alec Hammond <[email protected]>

Alec Hammond added 6 commits October 4, 2021 12:10

init

b00847b

get everything compiled and default passing tests

8481b07

new start

90472e4

anisotropic gradients are working

8a0369f

resolve meepgeom.cpp

3960357

fix multifrequency bug and other cleanup

b6d067c

smartalecH requested a review from ianwilliamson October 27, 2021 18:44

Alec Hammond added 3 commits October 27, 2021 16:10

add support for cyl coordindates

f8c2c93

Merge branch 'master' of https://github.com/NanoComp/meep into supixe…

d829397

…l_grad_mod

add rotated sapphire to adjoint solver test

afc1dab

ianwilliamson reviewed Oct 27, 2021

View reviewed changes

python/adjoint/optimization_problem.py Outdated Show resolved Hide resolved

python/adjoint/optimization_problem.py Outdated Show resolved Hide resolved

smartalecH mentioned this pull request Oct 28, 2021

Refactor wrapper and OptimizationProblem APIs to share code #1805

Open

get tests passing locally

1b635c2

Alec Hammond added 3 commits October 29, 2021 10:13

Merge branch 'master' of https://github.com/NanoComp/meep into supixe…

f6c0878

…l_grad_mod

enable dispersive material support

3de140a

some minor cleanup

64d7b25

smartalecH mentioned this pull request Nov 1, 2021

Better MaterialGrid symmetries test #1812

Closed

smartalecH requested a review from oskooi November 1, 2021 14:59

oskooi reviewed Nov 2, 2021

View reviewed changes

ianwilliamson reviewed Nov 2, 2021

View reviewed changes

clean up comments and readability issues and fix scaling of smoothing…

37174ef

… diameter

stevengj merged commit 8dd39f7 into NanoComp:master Nov 3, 2021

mochen4 mentioned this pull request Nov 9, 2021

Fix adjoint gradient with conductivities #1823

Closed

oskooi mentioned this pull request Nov 14, 2021

plot geometry for dispersive materials without initializing structure object #1827

Merged

smartalecH mentioned this pull request Nov 16, 2021

Subpixel smoothing gradients #1780

Merged

5 tasks

smartalecH mentioned this pull request Dec 6, 2021

Define what happens when β=∞ and u=η #1842

Merged

smartalecH mentioned this pull request Jan 7, 2022

compute adjoint gradient using restriction/interpolation only when anisotropic materials are present #1886

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Anisotropic material gradient support #1801

Anisotropic material gradient support #1801

smartalecH commented Oct 26, 2021 •

edited

Loading

smartalecH commented Oct 27, 2021

ianwilliamson commented Oct 27, 2021

smartalecH commented Oct 27, 2021 •

edited

Loading

ianwilliamson commented Oct 27, 2021

smartalecH commented Oct 27, 2021

ianwilliamson commented Oct 27, 2021

codecov-commenter commented Oct 28, 2021 •

edited

Loading

smartalecH commented Nov 1, 2021

oskooi Nov 2, 2021

smartalecH Nov 3, 2021

oskooi Nov 2, 2021

smartalecH Nov 2, 2021

smartalecH Nov 3, 2021

smartalecH Nov 3, 2021

oskooi Nov 2, 2021

smartalecH Nov 2, 2021

ianwilliamson Nov 2, 2021

smartalecH Nov 3, 2021

ianwilliamson Nov 2, 2021

smartalecH Nov 3, 2021

ianwilliamson Nov 2, 2021

smartalecH Nov 3, 2021

ianwilliamson Nov 2, 2021

smartalecH Nov 3, 2021

ianwilliamson Nov 2, 2021

smartalecH Nov 3, 2021

ianwilliamson Nov 2, 2021

smartalecH Nov 3, 2021

ianwilliamson Nov 2, 2021

smartalecH Nov 3, 2021

Anisotropic material gradient support #1801

Anisotropic material gradient support #1801

Conversation

smartalecH commented Oct 26, 2021 • edited Loading

smartalecH commented Oct 27, 2021

ianwilliamson commented Oct 27, 2021

smartalecH commented Oct 27, 2021 • edited Loading

ianwilliamson commented Oct 27, 2021

smartalecH commented Oct 27, 2021

ianwilliamson commented Oct 27, 2021

codecov-commenter commented Oct 28, 2021 • edited Loading

Codecov Report

smartalecH commented Nov 1, 2021

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

smartalecH commented Oct 26, 2021 •

edited

Loading

smartalecH commented Oct 27, 2021 •

edited

Loading

codecov-commenter commented Oct 28, 2021 •

edited

Loading