Luthaf / rascaline

Computing representations for atomistic machine learning
https://luthaf.fr/rascaline/
BSD 3-Clause "New" or "Revised" License
44 stars 13 forks source link

CGs: introduce general `TensorCorrelator` and refactor `DensityCorrelations` #316

Open jwa7 opened 2 months ago

jwa7 commented 2 months ago

Some changes:

  1. Add a keep_l_in_keys parameter for the DensityCorrelations calculator. This defaults to false, maintaining the previous behaviour. This makes it easier to write the pair correlator as l_{x} and k_{x} combination info needs to be kept in the keys for further CG tensor products. While this adds complexity to DensityCorrelations, the idea is for this complexity to in the end be wrapped by higher-level, more specific calculators (i.e. a lambda-SOAP calculator), such that the normal user doesn't need to worry about complex metadata things like this.
  2. body_order replaces correlation_order as the parameter in DensityCorrelations. "body_order" as a key dimension is returned in the output TensorMaps, as this is also useful to know for a general CG tensor product. Since we don't yet have global metadata, keeping it in keys is best for now.
  3. A match_keys argument has also been added, and the core routines that precompute tensor products of keys metadata have been modified. Internally, the keys metadata dimensions are categorised into 'standard' (i.e. order_nu, o3_lambda, o3_sigma), 'CG combination' (i..e the l and k lists), and 'other' key names. Standard and combination dimensions are handled as usual, but the other dimensions are either matched (based on match_keys argument) or fully multiplied.
  4. Arbitrary properties dimensions are now handled. Previously, the function assumes that (and handles by hard-coding) only "neighbor_type" and "n" properties dimensions are present in the TensorBlocks being combined. Now, arbitrary properties can be combined. The requirement is that the properties dimension names are different between the two blocks being combined. Property renaming is then handled by the wrapping method _correlate_density of DensityCorrelations. Properties of the arbitrary body order tensor are named (i.e. for radial channels) "n_1", "n2", ..., "n{nu - 1}", and the nu=1 tensor being combined on the current iteration is renamed so that its property dimension is "n{nu}". Suffixing with "{x}" is chosen and applied generically, except for dimensions suffixed with "_type", i.e. "neighbortype" becomes "neighbor{x}_type".

📚 Documentation preview 📚: https://rascaline--316.org.readthedocs.build/en/316/

github-actions[bot] commented 2 months ago

Here is a pre-built version of the code in this pull request: wheels.zip, you can install it locally by unzipping wheels.zip and using pip to install the file matching your system