Code structures for momentum-resolved phonon-loss simulations #184

TomaSusi · 2024-12-02T05:17:26Z

TomaSusi
Dec 2, 2024
Maintainer

Discussion of the ongoing development to enable simulations of momentum-resolved phonon energy loss in abTEM.

Code in branch phononloss:
https://github.com/abTEM/abTEM/tree/phononloss

TomaSusi · 2024-12-02T05:21:03Z

TomaSusi
Dec 2, 2024
Maintainer Author

To recap a private discussion with @pzeiger: to enable the required functionality, we need to implement a new EnergyAxis object and the related Measurement and potentially other classes.

When looking at the code, we realized there are a lot of classes in potentials/iam.py that are not specific to IAM potentials, and thus should be moved to their own module (called potentials/base.py). This way any changes we do to the data structures will not be dependent on what potential we use, and is in general a more modular code design.

0 replies

TomaSusi · 2024-12-02T05:25:59Z

TomaSusi
Dec 2, 2024
Maintainer Author

@pzeiger: I think I managed to do the basement surgery correctly – I moved non-IAM related classes to base.py, and I can import abTEM and it passes all the tests.

There's only only one thing I ended up unsure about: there was a validate_potential function in iam.py that instantiate the Potential class, and when I tried to move it to base.py, I thought to simply change that instantiation to the BasePotential class. However, BasePotential does not take a device argument, which Potential does (this is related to building the potential on CPU or GPU), and elsewhere in the code where validate_potential is called, it needs attributes that the pure BasePotential object does not have.

So I ended up leaving it in place at least for the time being.

2 replies

pzeiger Dec 3, 2024
Collaborator

I have had a brief look at the new base.py and there was duplicate code for the BaseField class. I pushed a commit where I removed the duplicate code.

TomaSusi Dec 4, 2024
Maintainer Author

Ah yes that was a simple omission, thanks for spotting.

TomaSusi · 2024-12-03T02:41:57Z

TomaSusi
Dec 3, 2024
Maintainer Author

@pzeiger: I took a stab at the missing Dask-related code for our new EnergyResolvedAtomsEnsemble object, and while I'm not sure I got everything right (changes pushed), the code isn't complaining about those for now. After creating the object, trying to build a potential out of it fails because I think our potentials/base.py module needs to be aware of the data type with the new EnergyAxis.

Quick code to test this:

import abtem
from abtem.inelastic.phonons import EnergyResolvedAtomsEnsemble

snaps = EnergyResolvedAtomsEnsemble(binned_snaps, energies)
potentials = abtem.Potential(snaps, slice_thickness = snaps[0].cell[2,2], sampling = sampling).build()

will result in the error:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[176], line 4
      1 energy = 200e3
      2 sampling = 0.05
----> 4 potentials = abtem.Potential(snaps, slice_thickness = snaps[0].cell[2,2], sampling = sampling).build()#.compute()

File [~/git/abTEM/abtem/potentials/base.py:544](http://localhost:8888/lab/workspaces/auto-s/tree/~/git/abTEM/abtem/potentials/base.py#line=543), in FieldBuilder.build(self, first_slice, last_slice, max_batch, lazy)
    541         for j, slic in enumerate(self.generate_slices(first_slice, last_slice)):
    542             array[j] = slic.array[0]
--> 544 output_potential = self._array_object(
    545     array,
    546     sampling=self._valid_sampling,
    547     slice_thickness=self.slice_thickness[first_slice:last_slice],
    548     exit_planes=self.exit_planes,
    549     ensemble_axes_metadata=self.ensemble_axes_metadata,
    550 )
    551 return output_potential

File [~/git/abTEM/abtem/potentials/base.py:885](http://localhost:8888/lab/workspaces/auto-s/tree/~/git/abTEM/abtem/potentials/base.py#line=884), in PotentialArray.__init__(self, array, slice_thickness, extent, sampling, exit_planes, ensemble_axes_metadata, metadata)
    882     metadata = {}
    883 metadata = {"label": "potential", "units": "eV [/](http://localhost:8888/) e", **metadata}
--> 885 super().__init__(
    886     array=array,
    887     slice_thickness=slice_thickness,
    888     extent=extent,
    889     sampling=sampling,
    890     exit_planes=exit_planes,
    891     ensemble_axes_metadata=ensemble_axes_metadata,
    892     metadata=metadata,
    893 )

File [~/git/abTEM/abtem/potentials/base.py:576](http://localhost:8888/lab/workspaces/auto-s/tree/~/git/abTEM/abtem/potentials/base.py#line=575), in FieldArray.__init__(self, array, slice_thickness, extent, sampling, exit_planes, ensemble_axes_metadata, metadata)
    571 self._exit_planes = _validate_exit_planes(
    572     exit_planes, len(self._slice_thickness)
    573 )
    574 self._grid = Grid(extent=extent, gpts=array.shape[-2:], sampling=sampling)
--> 576 super().__init__(
    577     array=array,
    578     ensemble_axes_metadata=ensemble_axes_metadata,
    579     metadata=metadata,
    580 )

File [~/git/abTEM/abtem/array.py:426](http://localhost:8888/lab/workspaces/auto-s/tree/~/git/abTEM/abtem/array.py#line=425), in ArrayObject.__init__(self, array, ensemble_axes_metadata, metadata, **kwargs)
    420 if len(array.shape) < self._base_dims:
    421     raise RuntimeError(
    422         f"{self.__class__.__name__} must be {self._base_dims}D or greater, not "
    423         f"{len(array.shape)}D"
    424     )
--> 426 self._check_axes_metadata()
    428 super().__init__(**kwargs)

File [~/git/abTEM/abtem/array.py:473](http://localhost:8888/lab/workspaces/auto-s/tree/~/git/abTEM/abtem/array.py#line=472), in ArrayObject._check_axes_metadata(self)
    471 def _check_axes_metadata(self) -> None:
    472     if len(self.shape) != len(self.axes_metadata):
--> 473         raise RuntimeError(
    474             f"number of array dimensions ({len(self.shape)}) does not match number"
    475             f"of axis metadata items ({len(self.axes_metadata)})"
    476         )
    478     for n, axis in zip(self.shape, self.axes_metadata):
    479         if isinstance(axis, OrdinalAxis) and len(axis) != n:

RuntimeError: number of array dimensions (4) does not match numberof axis metadata items (5)

I won't have time to work on that further today, so if you can take a crack at it you can use any list of Atoms for the binned_snaps object for testing.

14 replies

pzeiger Dec 10, 2024
Collaborator

Interesting. I have turned of the autoreload of modules and restarted Jupyter and now I also get your error. Before that it was always and the error

ValueError: array module specification dask.array<_wrap_build_potential, shape=(2, 5, 1, 345, 341), dtype=float32, chunksize=(1, 1, 1, 345, 341), chunktype=numpy.ndarray> not recognized

I can also take a look at it later today.

TomaSusi Dec 11, 2024
Maintainer Author

First thing I encountered was a realization that while our new potential could be built, it couldn't actually be computed:

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Cell In[206], line 4
      1 energy = 200e3
      2 sampling = 0.05
----> 4 potentials = abtem.Potential(snaps, slice_thickness = snaps[0].cell[2,2], sampling = sampling).build().compute()

[omitting irrelevant parts of the stacktrace]

File [~/git/abTEM/abtem/potentials/iam.py:320](http://localhost:8888/lab/workspaces/auto-s/tree/~/git/abTEM/abtem/potentials/iam.py#line=319), in _FieldBuilderFromAtoms._from_partitioned_args_func(cls, frozen_phonons_partial, *args, **kwargs)
    316 @classmethod
    317 def _from_partitioned_args_func(cls, *args, frozen_phonons_partial, **kwargs):
    318     args = unpack_blockwise_args(args)
--> 320     frozen_phonons = frozen_phonons_partial(*args)
    321     frozen_phonons = frozen_phonons.item()
    323     new_potential = cls(frozen_phonons, **kwargs)

File [~/git/abTEM/abtem/inelastic/phonons.py:910](http://localhost:8888/lab/workspaces/auto-s/tree/~/git/abTEM/abtem/inelastic/phonons.py#line=909), in EnergyResolvedAtomsEnsemble._from_partition_args_func(*args, **kwargs)
    908 args = unpack_blockwise_args(args)
    909 energy_resolved_snapshots = args[0]
--> 910 energies = args[1]
    911 snapshots_stack = EnergyResolvedAtomsEnsemble(energy_resolved_snapshots, energies, **kwargs)
    912 return _wrap_with_array(snapshots_stack, 2)

IndexError: tuple index out of range

In fact, the potential object being built in iam.py is not aware of the energies and thus is not passing those back to the function of our new object. This leads back to the function _from_partitioned_args_func of the _FieldBuilderFromAtoms class, which would also need to understand that the potential it is asked to calculate has this new energy structure.

However, that I think is not going to be a viable path, as we would again have to reimplement a lot of the functionality that is already there for each frozen-phonon member of the energy ensemble. So I think we need to tackle this one level up, and modify the potential object to be aware of the energies, and when called handles the iterations over each energy of the ensemble.

I think this would need to be in base.py, and we would need to modify both BasePotential and PotentialArray.

pzeiger Dec 11, 2024
Collaborator

I worked today on this and pushed some very incremental progress. I have reached a point where a NotImplementedError() is raised with the implementation EnergyResolvedAtomsEnsemble2. I think the error is about that splitting of the arrays along more than 1 ensemble axis is not implemented. I will have a look at that in more detail tomorrow as it would be also a problem for the implementation EnergyResolvedAtomsEnsemble. I tried to bring that one to a similar state, but struggled with actually getting the chunking right and assembling the proper dask array. I noticed that the chunking we had made was not proper for how we use the args argument in the method _from_partition_args_func(), but I have not found a way to build the right dask array yet.

TomaSusi Dec 12, 2024
Maintainer Author

Hmm... I'm not sure what you mean by "arrays" being split, but objects with more than one ensemble axis definitely can be handled, e.g. frozen phonons with a defocus ensemble. Personally I don't think the problem is in chunking (except in the sense that each FrozenPhonon member of the energy ensemble can be chunked separately). The problem I see is that the potential computation does not include the energy dimension as something that needs to be handled differently to frozen phonon axes.

Maybe we could have a call morning Thursday your time to discuss a bit more? I'm free around 9:30 CET.

pzeiger Dec 12, 2024
Collaborator

Yes, that would work. I will be online at 9:30 CET.

TomaSusi · 2024-12-16T01:04:01Z

TomaSusi
Dec 16, 2024
Maintainer Author

I was mulling over our current approach, and what we discussed at the end of last week: all of the existing AtomsEnsemble, Potential, and multislice machinery really don't need to know anything about the energy axis, and thus it actually isn't necessary to change any of that to get a usable workflow.

So went back to the original approach of using stacks, and I now have a workflow that just stacks stuff on top of each other, retaining the FrozenPhonon snapshops optionally but doing an ensemble mean by default, and resulting in an diffraction intensities stack that has I_vib, I_coh and I_incoh as the first ensemble axis and the energy bins as the second – with all of the abTEM measurement structure retained up until that point, including the ability to save this as a single .zarr file.

The only thing that is missing is the ability to apply a detector on this, which I think is because the MeasurementsEnsemble object doesn't understand there can be an extra ensemble axis – the error below is only happening for the visualization, the actual detected values get calculated correctly.

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Cell In[197], line 1
----> 1 meas.show()

File [~/git/abTEM/abtem/measurements.py:838](http://localhost:8888/lab/workspaces/auto-s/tree/~/git/abTEM/abtem/measurements.py#line=837), in MeasurementsEnsemble.show(self, type, ax, power, common_scale, explode, overlay, figsize, title, units, interact, display, **kwargs)
    822 # raise RuntimeError("Cannot infer pixel area from metadata.")
    823 
    824 # if display_axes != (-2, -1):
   (...)
    833 #     if i not in display_axes
    834 # ]
    836 artist_type = LinesArtist
--> 838 visualization = Visualization(
    839     measurement=self,
    840     ax=ax,
    841     artist_type=artist_type,
    842     # power=power,
    843     aspect=False,
    844     share_x=True,
    845     share_y=common_scale,
    846     common_scale=common_scale,
    847     explode=explode,
    848     overlay=overlay,
    849     figsize=figsize,
    850     # interact=interact,
    851     title=title,
    852     **kwargs,
    853 )
    855 if common_scale is False and visualization._explode:
    856     visualization.axes.set_sizes(padding=0.8)

File [~/git/abTEM/abtem/visualize/visualizations.py:269](http://localhost:8888/lab/workspaces/auto-s/tree/~/git/abTEM/abtem/visualize/visualizations.py#line=268), in Visualization.__init__(self, measurement, ax, artist_type, figsize, aspect, common_scale, value_limits, overlay, explode, share_x, share_y, cbar, interactive, title, xlim, ylim, convert_complex, **kwargs)
    262     row_titles = [
    263         l.format_title(".3g", include_label=i == 0)
    264         for i, l in enumerate(axes_metadata)
    265     ]
    267     self.set_row_titles(row_titles)
--> 269 self._make_new_artists(artist_type=self._artist_type, **kwargs)
    271 self.adjust_coordinate_limits_to_artists(xlim=xlim, ylim=ylim)
    273 if common_scale:

File [~/git/abTEM/abtem/visualize/visualizations.py:598](http://localhost:8888/lab/workspaces/auto-s/tree/~/git/abTEM/abtem/visualize/visualizations.py#line=597), in Visualization._make_new_artists(self, artist_type, **kwargs)
    593     # else:
    594     #     caxes = [_make_cax(ax) for _ in range(artist_type.num_cbars)]
    596     measurement = self._reduce_measurement(self._indices, i)
--> 598     artist = artist_type(
    599         ax=ax,
    600         caxes=caxes,
    601         measurement=measurement,
    602         **kwargs,
    603     )
    605     itemset(artists, i, artist)
    607 self._artists = artists

File [~/git/abTEM/abtem/visualize/artists.py:202](http://localhost:8888/lab/workspaces/auto-s/tree/~/git/abTEM/abtem/visualize/artists.py#line=201), in LinesArtist.__init__(self, ax, measurement, caxes, label, units, legend, **kwargs)
    198 super().__init__(ax=ax, measurement=measurement)
    200 y = self._reshape_data(measurement.array)
--> 202 x = measurement.axes_metadata[-1].coordinates(measurement.shape[-1])
    204 if label is None and measurement.ensemble_shape:
    205     labels = []

IndexError: list index out of range

So I think we can tackle this much simpler by either modifying the existing MeasurementsArray object or creating a new one that understands the energy ensemble, and if we want to be elegant about it, doing that with an EnergyAxis earlier in the workflow. But this feels like a much smaller task to accomplish, and I think we can simply leave the Dask plumbing as it is.

I'd be available to call again to show you this approach and discuss the next steps around the usual time :)

1 reply

TomaSusi Dec 20, 2024
Maintainer Author

I took a whack at implementing a new module spectra.py (instead of adding these to the already overloaded measurements.py, and implementing the base spectrum objects as well as the MomentumResolvedSpectrum class.

These can be instantiated from the old data structures as follows:

from abtem.spectra import MomentumResolvedSpectrum

Eqspectrum = MomentumResolvedSpectrum.from_array_and_metadata(array = linex.array, axes_metadata = linex.ensemble_axes_metadata)

Here we see one architecting issue that I ran into immediately. For regular measurement objects, the base axes (axes_metadata) describe the contents of the arrays, while the things that are stacked (which I how I created my objects) are described by ensemble_axes_metadata. I tried to "cheat" and use the energy and momentum axes from my measurement ensemble as the axes, but this runs into a problem when trying we try to instantiate the underlying ArrayObject class:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[154], line 9
      5 del modules["abtem.array"]
      7 from abtem.spectra import MomentumResolvedSpectrum
----> 9 Eqspectrum = MomentumResolvedSpectrum.from_array_and_metadata(array = linex.array, axes_metadata = linex.ensemble_axes_metadata)

File [~/git/abTEM/abtem/spectra.py:529](http://localhost:8888/lab/workspaces/auto-s/tree/~/git/abTEM/abtem/spectra.py#line=528), in MomentumResolvedSpectrum.from_array_and_metadata(cls, array, axes_metadata, metadata)
    526 else:
    527     raise RuntimeError()
--> 529 return cls(
    530     array,
    531     values=values,
    532     axes_metadata=axes_metadata[:-2],
    533     metadata=metadata,
    534 )

File [~/git/abTEM/abtem/spectra.py:490](http://localhost:8888/lab/workspaces/auto-s/tree/~/git/abTEM/abtem/spectra.py#line=489), in MomentumResolvedSpectrum.__init__(self, array, values, axes_metadata, metadata)
    481 def __init__(
    482     self,
    483     array: da.core.Array | np.array,
   (...)
    486     metadata: Optional[Dict] = None,
    487 ):
    488     self._values = values
--> 490     super().__init__(
    491         array=array,
    492         axes_metadata=axes_metadata,
    493         metadata=metadata,
    494     )

File [~/git/abTEM/abtem/spectra.py:110](http://localhost:8888/lab/workspaces/auto-s/tree/~/git/abTEM/abtem/spectra.py#line=109), in BaseSpectra.__init__(self, array, axes_metadata, metadata)
    104 def __init__(
    105     self,
    106     array: np.ndarray | da.core.Array,
    107     axes_metadata: list[AxisMetadata],
    108     metadata: dict,
    109 ):
--> 110     super().__init__(
    111         array=array,
    112         ensemble_axes_metadata=axes_metadata,
    113         metadata=metadata,
    114     )

File [~/git/abTEM/abtem/array.py:428](http://localhost:8888/lab/workspaces/auto-s/tree/~/git/abTEM/abtem/array.py#line=427), in ArrayObject.__init__(self, array, ensemble_axes_metadata, metadata, **kwargs)
    422 if len(array.shape) < self._base_dims:
    423     raise RuntimeError(
    424         f"{self.__class__.__name__} must be {self._base_dims}D or greater, not "
    425         f"{len(array.shape)}D"
    426     )
--> 428 self._check_axes_metadata()
    430 super().__init__(**kwargs)

File [~/git/abTEM/abtem/array.py:484](http://localhost:8888/lab/workspaces/auto-s/tree/~/git/abTEM/abtem/array.py#line=483), in ArrayObject._check_axes_metadata(self)
    480 for n, axis in zip(self.shape, self.axes_metadata):
    481     #if isinstance(axis, EnergyAxis) or isinstance(axis, MomentumAxis):
    482     #    pass
    483     #elif isinstance(axis, OrdinalAxis) and len(axis) != n:
    484     if isinstance(axis, OrdinalAxis) and len(axis) != n:
--> 485         raise RuntimeError(
    487             f"number of values for ordinal axis ({len(axis)}), does not match "
    488             f"size of dimension ({n})."
    488         )

RuntimeError: number of values for ordinal axis (51), does not match size of dimension (57).

This check is triggered because we based our EnergyAxis and MomentumAxis objects on OrdinalAxis, which is used to describe stacks, and which triggers this error.

You can see some code that's commented out, which checks if the ordinal axis is in fact one of our new axes and bypasses this check, but this only leads to a further problem down the line where another axis metadata check adds the base axes and ordinal axes together and gets the wrong number of axes for the 2D array:

File [~/git/abTEM/abtem/array.py:428](http://localhost:8888/lab/workspaces/auto-s/tree/~/git/abTEM/abtem/array.py#line=427), in ArrayObject.__init__(self, array, ensemble_axes_metadata, metadata, **kwargs)
    422 if len(array.shape) < self._base_dims:
    423     raise RuntimeError(
    424         f"{self.__class__.__name__} must be {self._base_dims}D or greater, not "
    425         f"{len(array.shape)}D"
    426     )
--> 428 self._check_axes_metadata()
    430 super().__init__(**kwargs)

File [~/git/abTEM/abtem/array.py:475](http://localhost:8888/lab/workspaces/auto-s/tree/~/git/abTEM/abtem/array.py#line=474), in ArrayObject._check_axes_metadata(self)
    473 def _check_axes_metadata(self) -> None:
    474     if len(self.shape) != len(self.axes_metadata):
--> 475         raise RuntimeError(
    476             f"number of array dimensions ({len(self.shape)}) does not match number"
    477             f"of axis metadata items ({len(self.axes_metadata)})"
    478         )
    480     for n, axis in zip(self.shape, self.axes_metadata):
    481         if isinstance(axis, EnergyAxis) or isinstance(axis, MomentumAxis):

RuntimeError: number of array dimensions (2) does not match number of axis metadata items (4).

One solution would be of course to base our energy and momentum axes on LinearAxis instead (note that NonLinearAxis would not help since that is a subclass or OrdinalAxis), but then stacking would no longer work... argh.

Despite a remarkably modular abTEM design, some fundamental distinctions between base and ensemble axes are really hard to bend. Anyways, I pushed the spectra.py but not the attempted hacks to array.py.

Code structures for momentum-resolved phonon-loss simulations #184

Uh oh!

Uh oh!

TomaSusi Dec 2, 2024 Maintainer

Replies: 4 comments · 17 replies

Uh oh!

TomaSusi Dec 2, 2024 Maintainer Author

Uh oh!

TomaSusi Dec 2, 2024 Maintainer Author

Uh oh!

Uh oh!

pzeiger Dec 3, 2024 Collaborator

Uh oh!

TomaSusi Dec 4, 2024 Maintainer Author

Uh oh!

TomaSusi Dec 3, 2024 Maintainer Author

Uh oh!

pzeiger Dec 10, 2024 Collaborator

Uh oh!

Uh oh!

TomaSusi Dec 11, 2024 Maintainer Author

Uh oh!

pzeiger Dec 11, 2024 Collaborator

Uh oh!

TomaSusi Dec 12, 2024 Maintainer Author

Uh oh!

pzeiger Dec 12, 2024 Collaborator

Uh oh!

TomaSusi Dec 16, 2024 Maintainer Author

Uh oh!

TomaSusi Dec 20, 2024 Maintainer Author

TomaSusi
Dec 2, 2024
Maintainer

Replies: 4 comments 17 replies

TomaSusi
Dec 2, 2024
Maintainer Author

TomaSusi
Dec 2, 2024
Maintainer Author

pzeiger Dec 3, 2024
Collaborator

TomaSusi Dec 4, 2024
Maintainer Author

TomaSusi
Dec 3, 2024
Maintainer Author

pzeiger Dec 10, 2024
Collaborator

TomaSusi Dec 11, 2024
Maintainer Author

pzeiger Dec 11, 2024
Collaborator

TomaSusi Dec 12, 2024
Maintainer Author

pzeiger Dec 12, 2024
Collaborator

TomaSusi
Dec 16, 2024
Maintainer Author

TomaSusi Dec 20, 2024
Maintainer Author